bug fixes

Andrewzh112 · Dec 2, 2020 · a32bc9c · a32bc9c
1 parent 0ac2d0d
commit a32bc9c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/simsiam/main.py b/simsiam/main.py
@@ -122,7 +122,7 @@ def cosine_loss(p, z):
         linear_classifier.fit(feature_bank, feature_labels)
 
         y_preds, y_trues = [], []
-        for data, target in feature_loader:
+        for data, target in test_loader:
             data = data.to(device)
             with torch.no_grad():
                 feature = model(data, istrain=False)