explicitly add weights_only=False to torch.load calls

Nerogar · Oct 13, 2024 · fa78a18 · fa78a18
1 parent 9ca340e
commit fa78a18
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/mgds/pipelineModules/DiskCache.py b/src/mgds/pipelineModules/DiskCache.py
@@ -218,7 +218,7 @@ def fn(group_index, in_index, in_variation):
 
                 if self.aggregate_cache[group_key][in_variation] is None:
                     self.aggregate_cache[group_key][in_variation] = \
-                        torch.load(os.path.realpath(os.path.join(cache_dir, 'aggregate.pt')))
+                        torch.load(os.path.realpath(os.path.join(cache_dir, 'aggregate.pt')), weights_only=False)
 
     def __get_input_index(self, out_variation: int, out_index: int) -> (str, int, int):
         offset = 0
@@ -251,7 +251,7 @@ def get_item(self, index: int, requested_name: str = None) -> dict:
 
         elif requested_name in self.split_names:
             cache_path = os.path.join(self.__get_cache_dir(group_key, in_variation), str(group_index) + '.pt')
-            split_item = torch.load(os.path.realpath(cache_path))
+            split_item = torch.load(os.path.realpath(cache_path), weights_only=False)
 
             for name in self.split_names:
                 item[name] = split_item[name]