Move to cuda unconditionally so pp-only run works

pytorch · Feb 9, 2024 · b382345 · b382345
1 parent e1b61c3
commit b382345
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchtrain/parallelisms/parallelize_llama.py b/torchtrain/parallelisms/parallelize_llama.py
@@ -185,4 +185,6 @@ def parallelize_llama(model, world_mesh, parallel_dims, args):
 
         rank0_log("Applied FSDP to the model...")
 
+    # redundant if FSDP is used, but ensure the model is on device consistently regardless with parallelisms were used
+    model.cuda()
     return model