Update modeling_quiet.py

modeling_quiet.py

@@ -1567,15 +1567,16 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 				return x
 			return x.repeat_interleave(n, dim=0)
 
-		if self.n_passes > 1:
+		if self.n_passes > 1 and input_ids is not None:
 			input_ids = none_repeat_interleave(input_ids, self.n_passes)
-			attention_mask = none_repeat_interleave(attention_mask, self.n_passes)
-			position_ids = none_repeat_interleave(position_ids, self.n_passes)
-			inputs_embeds = none_repeat_interleave(inputs_embeds, self.n_passes)
-			labels = none_repeat_interleave(labels, self.n_passes)
+			attention_mask = none_repeat_interleave(attention_mask, self.n_passes) if attention_mask is not None else None
+			position_ids = none_repeat_interleave(position_ids, self.n_passes) if position_ids is not None else None
+			inputs_embeds = none_repeat_interleave(inputs_embeds, self.n_passes) if inputs_embeds is not None else None
+			labels = none_repeat_interleave(labels, self.n_passes) if labels is not None else None
 			if past_key_values is not None:
 				past_key_values = [none_repeat_interleave(p, self.n_passes) for p in past_key_values]
-		cur_token_indices = torch.arange(input_ids.shape[1], device=input_ids.device)
+		if input_ids is not None:
+			cur_token_indices = torch.arange(input_ids.shape[1], device=input_ids.device)
 
 		self.tokenizer_has_start_thought_token = True
 		self.tokenizer_has_end_thought_token = True