1. Fixed device assignemnt.

Signed-off-by: Micha Livne <[email protected]>
NVIDIA · Mar 6, 2023 · 0d51dfa · 0d51dfa
1 parent 4fe70c0
commit 0d51dfa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_lm_encoder_decoder_model.py b/nemo/collections/nlp/models/language_modeling/megatron_lm_encoder_decoder_model.py
@@ -1255,11 +1255,11 @@ def dummy():
             else:
                 predicted_log_probs = torch.zeros(
                     (predicted_log_probs.shape[0], predicted_log_probs.shape[1]), dtype=self.autocast_dtype
-                ).to(device)
+                ).cuda()
                 predicted_tokens_dec = torch.zeros(
                     (predicted_tokens_dec.shape[0], predicted_tokens_dec.shape[1] + 1),
                     dtype=predicted_tokens_dec.dtype,
-                ).to(device)
+                ).cuda()
 
             if self.cfg.get('pipeline_model_parallel_size', 1) > 1:
                 # Broadcast from the last pipeline stage to all other model-parallel ranks.