custom input image size fixed.

ghimiredhikura · Dec 6, 2019 · bcad024 · bcad024
1 parent 994b264
commit bcad024
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -34,7 +34,7 @@
     class_names = load_classes("data/classes.names")
 
     # Initiate model
-    model = Darknet(opt.model_def).to(device)
+    model = Darknet(opt.model_def, img_size=opt.img_size).to(device)
     model.apply(weights_init_normal)
 
     # If specified we start from checkpoint