From d566677b34762951c05d43f011f60be7ee2c594e Mon Sep 17 00:00:00 2001
From: HobbitLong <yonglong@mit.edu>
Date: Wed, 10 Jun 2020 23:00:36 -0400
Subject: [PATCH] fix

---
 main_supcon.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/main_supcon.py b/main_supcon.py
index 8467b189..8688667a 100644
--- a/main_supcon.py
+++ b/main_supcon.py
@@ -190,6 +190,7 @@ def train(train_loader, model, criterion, optimizer, epoch, opt):
         data_time.update(time.time() - end)
 
         images = torch.cat([images[0], images[1]], dim=0)
+        images = images.cuda(non_blocking=True)
         labels = labels.cuda(non_blocking=True)
         bsz = labels.shape[0]