use tempfile to avoid name conflict

geyang · Sep 18, 2021 · e7f754e · e7f754e
1 parent e184579
commit e7f754e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 5 deletions.
diff --git a/ml_logger/VERSION b/ml_logger/VERSION
@@ -1 +1 @@
-0.8.32
+0.8.33
diff --git a/ml_logger/ml_logger/ml_logger.py b/ml_logger/ml_logger/ml_logger.py
@@ -1605,14 +1605,16 @@ def download_file(self, *keys, path=None, to, relative=False):
             f.write(buf.getbuffer())
 
     def load_torch(self, *keys, path=None, map_location=None, **kwargs):
-        import torch
+        import torch, tempfile
         path = pJoin(*keys, path)
         if path.lower().startswith('s3://'):
-            fn_or_buff = os.path.basename(path)
-            self.download_s3(path[5:], to=fn_or_buff)
+            postfix = os.path.basename(path)
+            with tempfile.NamedTemporaryFile(suffix=f'.{postfix}') as ntp:
+                self.download_s3(path[5:], to=ntp.name)
+                return torch.load(ntp, map_location=map_location, **kwargs)
         else:
             fn_or_buff = self.load_file(path)
-        return torch.load(fn_or_buff, map_location=map_location, **kwargs)
+            return torch.load(fn_or_buff, map_location=map_location, **kwargs)
 
     torch_load = load_torch