Fix PipelineEngine.eval_batch result

microsoft · Apr 20, 2023 · 7a632c6 · 7a632c6
1 parent 4de4d2a
commit 7a632c6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deepspeed/runtime/pipe/engine.py b/deepspeed/runtime/pipe/engine.py
@@ -500,7 +500,7 @@ def _bcast_pipe_scalar(self, data, src_rank=None, dtype=torch.float32):
         assert src_rank in self.grid.pp_group
 
         if self.global_rank == src_rank:
-            result = data.clone().detach()
+            result = data.clone().detach().type(dtype).to(self.device)
         else:
             result = torch.Tensor([0.]).type(dtype).to(self.device)