fix

os1a · Jun 11, 2020 · d566677 · d566677
1 parent 6145bc3
commit d566677
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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]