Update DMI.py

Newbeeer · Sep 30, 2019 · 95d7d46 · 95d7d46
1 parent cc74eb8
commit 95d7d46
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/dogcat/DMI.py b/dogcat/DMI.py
@@ -24,7 +24,7 @@ def train(train_loader, model, optimizer, criterion=CE):
 
     for i, (idx, input, target) in enumerate(tqdm(train_loader)):
         if idx.size(0) != batch_size:
-            break
+            continue
 
         input = torch.autograd.Variable(input.cuda())
         target = torch.autograd.Variable(target.cuda())
@@ -67,7 +67,7 @@ def validate(valid_loader, model, criterion):
 
         for i, (idx, input, target) in enumerate(tqdm(valid_loader)):
             if i == 1:
-                continue
+                break
 
             input = torch.autograd.Variable(input.cuda())
             target = torch.autograd.Variable(target.cuda())
@@ -109,4 +109,4 @@ def evaluate(path):
 
 if __name__ == '__main__':
     main_dmi()
-    evaluate('./dog_model_dmi_' + str(args.r) + '_' + str(args.s))
+    evaluate('./dog_model_dmi_' + str(args.r) + '_' + str(args.s))