Update metadata.yaml

pytorch · Oct 12, 2023 · 2214663 · 2214663
1 parent 7eb245a
commit 2214663
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/torchbenchmark/models/squeezenet1_1/metadata.yaml b/torchbenchmark/models/squeezenet1_1/metadata.yaml
@@ -5,6 +5,9 @@ eval_benchmark: true
 eval_deterministic: false
 eval_nograd: true
 not_implemented:
-- test: train
+  - device: cpu
+    test: train
+  - device: cuda
+    test: train
 train_benchmark: true
 train_deterministic: false