Update modeling_quiet.py

modeling_quiet.py

@@ -1847,8 +1847,7 @@ class QuietForCausalLM(QuietPreTrainedModel):
                             [shift_labels, padding],
                             dim=-1
                         )
-                        # 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).to(probabilities_2d.device)
+                        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
                     else:
                         continue