Update modeling_quiet.py

modeling_quiet.py

@@ -1236,6 +1236,21 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         if input_ids.dim() == 1:
             input_ids = input_ids.unsqueeze(0)
             attention_mask = attention_mask.unsqueeze(0) if attention_mask is not None else None
+            position_ids = position_ids.unsqueeze(0) if position_ids is not None else None
+    
+        seq_len = input_ids.shape[1]
+    
+        if position_ids is None:
+            position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+        else:
+            # Handle the case when position_ids is an empty tensor
+            if position_ids.numel() == 0:
+                position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+                position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
+            else:
+                position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
+
             
         output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
         output_hidden_states = (