Skip to content

Commit

Permalink
Rearrange export input after checks (ultralytics#1118)
Browse files Browse the repository at this point in the history
img size checks are warnings rather than errors, so current implementation allows improperly formed model inputs.
  • Loading branch information
glenn-jocher committed Oct 11, 2020
1 parent 899fe42 commit 5bad741
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions models/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,6 @@
set_logging()
t = time.time()

# Input
img = torch.zeros((opt.batch_size, 3, *opt.img_size)) # image size(1,3,320,192) iDetection

# Load PyTorch model
model = attempt_load(opt.weights, map_location=torch.device('cpu')) # load FP32 model
labels = model.names
Expand All @@ -40,6 +37,9 @@
gs = int(max(model.stride)) # grid size (max stride)
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples

# Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size) # image size(1,3,320,192) iDetection

# Update model
for k, m in model.named_modules():
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
Expand Down

0 comments on commit 5bad741

Please sign in to comment.