Update modeling_quiet.py

modeling_quiet.py

@@ -1847,8 +1847,11 @@ class QuietForCausalLM(QuietPreTrainedModel):
                             [shift_labels, padding],
                             dim=-1
                         )
-                        probabilities_2d = F.one_hot(new_rm_tokens, num_classes=self.vocab_size).reshape(-1, self.vocab_size).to(probabilities_2d.dtype).to(probabilities_2d.device)
-                        skip_sampling = True
+						# Before converting rm tokens to one-hot, clamp the values to ensure they are within the valid index range
+						new_rm_tokens = torch.clamp(new_rm_tokens, 0, self.vocab_size - 1)
+
+						# Now safely convert rm tokens to one-hot
+						probabilities_2d = F.one_hot(new_rm_tokens, num_classes=self.vocab_size).reshape(-1, self.vocab_size).to(probabilities_2d.dtype)
                     else:
                         continue
                 temperature = self.gumbel_temperature if self.training else 0.001