Move to cuda unconditionally so pp-only run works (#50)

pytorch · Feb 10, 2024 · 1a0b9fd · 1a0b9fd
1 parent e1b61c3
commit 1a0b9fd
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 enabled, but ensure the model is on device regardless of which parallelisms were used
+    model.cuda()
     return model