Update weights_only due to change in default in torch>=2.6

deepspeedai · Feb 10, 2025 · 196d624 · 196d624
1 parent b90ffab
commit 196d624
Showing 1 changed file with 1 addition 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,7 @@ 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: