Update weights_only due to change in default in torch 2.6+ (#957)

* Update weights_only due to change in default in torch>=2.6 Signed-off-by: Logan Adams <[email protected]> * formatting Signed-off-by: Logan Adams <[email protected]> --------- Signed-off-by: Logan Adams <[email protected]>
deepspeedai · Feb 12, 2025 · 83757d9 · 83757d9
1 parent b90ffab
commit 83757d9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/applications/DeepSpeed-Chat/dschat/utils/data/data_utils.py b/applications/DeepSpeed-Chat/dschat/utils/data/data_utils.py
@@ -375,7 +375,9 @@ def create_prompt_dataset(local_rank,
         torch.save(train_dataset, train_fname)
         torch.save(eval_dataset, eval_fname)
     torch.distributed.barrier()
-    return torch.load(train_fname), torch.load(eval_fname)
+    return torch.load(train_fname,
+                      weights_only=False), torch.load(eval_fname,
+                                                      weights_only=False)
 
 
 class DataCollatorReward: