Home Explore Help
Register Sign In
marx
/
llama-recipes
1
0
Fork 0
Files Issues 0 Pull Requests 0 Wiki
Browse Source

fix #90

lchu 1 year ago
parent
c19c5c69aa
commit
0c51b47262
1 changed files with 2 additions and 1 deletions
Split View Show Diff Stats
  1. 2 1
      llama_finetuning.py

+ 2 - 1
llama_finetuning.py
View File 

@@ -86,7 +86,8 @@ def main(**kwargs):
 
         world_size = int(os.environ["WORLD_SIZE"])
 
 
     if torch.distributed.is_initialized():
 
-        torch.cuda.set_device(rank)
 
+        torch.cuda.set_device(local_rank)
 
+        clear_gpu_cache(local_rank)
 
         setup_environ_flags(rank)
 
 
     # Calculate gradient accumulation steps