Update references/classification/train.py

pytorch · Dec 14, 2023 · b88457e · b88457e
1 parent 2255c4a
commit b88457e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/references/classification/train.py b/references/classification/train.py
@@ -127,7 +127,8 @@ def load_data(traindir, valdir, args):
     if args.cache_dataset and os.path.exists(cache_path):
         # Attention, as the transforms are also cached!
         print(f"Loading dataset_train from {cache_path}")
-        dataset, _ = torch.load(cache_path, weights_only=True)
+        # TODO: this could probably be weights_only=True
+        dataset, _ = torch.load(cache_path, weights_only=False)
     else:
         # We need a default value for the variables below because args may come
         # from train_quantization.py which doesn't define them.