specify torch.load weights_only

mit-ccc · Sep 24, 2024 · 361291d · 361291d
1 parent 17e1f0a
commit 361291d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/images/sentiment-topic/topic-embeds.py b/images/sentiment-topic/topic-embeds.py
@@ -122,7 +122,7 @@ def _combine_cached_values(self):
             ids += [pd.read_csv(fname + '.csv')]
 
             with open(fname + '.pt', 'rb') as obj:
-                embeds += [torch.load(obj, 'cpu')]
+                embeds += [torch.load(obj, 'cpu', weights_only=True)]
         ids = pd.concat(ids, axis=0)
         embeds = torch.cat(embeds, dim=0)
 

diff --git a/images/sentiment-topic/topic-umap.py b/images/sentiment-topic/topic-umap.py
@@ -60,7 +60,10 @@
 
     embeds_file = os.path.join(data_dir, 'comment-topics/sentence-embeds.pt')
     with open(embeds_file, 'rb') as obj:
-        embeds = torch.load(obj, 'cpu').float().numpy()[train_mask, ...]
+        embeds = torch.load(obj, 'cpu', weights_only=True) \
+            .float() \
+            .numpy() \
+            [train_mask, ...]
 
     params = {
         'n_neighbors': 15,