fix remote file naming (#3173)

mosaicml · Apr 3, 2024 · 5b04b00 · 5b04b00
1 parent d7f8a09
commit 5b04b00
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/composer/callbacks/memory_snapshot.py b/composer/callbacks/memory_snapshot.py
@@ -76,7 +76,7 @@ def __init__(
         max_entries: int = 100000,
         folder: str = '{run_name}/torch_traces',
         filename: str = 'rank{rank}.{batch}.memory_snapshot',
-        remote_file_name: Optional[str] = '{run_name}/torch_memory_traces/rank{rank}.{batch}.memory_snapshot',
+        remote_file_name: Optional[str] = '{run_name}/torch_memory_traces',
         overwrite: bool = False,
     ) -> None:
         self.batches_left_to_skip = skip_batches
@@ -179,7 +179,7 @@ def export_memory_snapshot(self, state: State, logger: Logger) -> None:
 
             if self.remote_path_in_bucket is not None:
                 for f in [snapshot_file, trace_plot_file]:
-                    remote_file_name = (self.remote_path_in_bucket + os.path.basename(f)).lstrip('/')
+                    remote_file_name = os.path.join(self.remote_path_in_bucket, os.path.basename(f)).lstrip('/')
                     log.info(f'Uploading memory snapshot to remote: {remote_file_name} from {f}')
                     try:
                         logger.upload_file(remote_file_name=remote_file_name, file_path=f, overwrite=self.overwrite)