Specify weights_only=False to avoid a FutureWarning

dagardner-nv · Oct 15, 2024 · 10ad08a · 10ad08a
1 parent be4597a
commit 10ad08a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/gnn_fraud_detection_pipeline/stages/model.py b/examples/gnn_fraud_detection_pipeline/stages/model.py
@@ -407,7 +407,7 @@ def load_model(model_dir: str,
                       n_layers=hyperparameters['n_layers'],
                       embedding_size=hyperparameters['embedding_size'],
                       target=hyperparameters['target_node']).to(device)
-    model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pt')))
+    model.load_state_dict(torch.load(os.path.join(model_dir, 'model.pt'), weights_only=False))
 
     return model, graph, hyperparameters