adding PT cuda allocation expand flag

utils/train_utils.py

@@ -253,6 +253,7 @@ def setup_environ_flags(rank):
     os.environ["TORCH_SHOW_CPP_STACKTRACES"] = str(1)
     os.environ["NCCL_ASYNC_ERROR_HANDLING"] = str(1)
     os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL"
+    os.environ['PYTORCH_CUDA_ALLOC_CONF']='expandable_segments:True'
     if rank == 0:
         print(f"--> Running with torch dist debug set to detail")