Skip to content

Commit

Permalink
Fix PyTorch Hub export inference shapes (#6949)
Browse files Browse the repository at this point in the history
  • Loading branch information
MichaelAnderson-AI committed Mar 11, 2022
1 parent 90486e2 commit 13051d9
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions models/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -544,10 +544,9 @@ def forward(self, imgs, size=640, augment=False, profile=False):
g = (size / max(s)) # gain
shape1.append([y * g for y in s])
imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
shape1 = [make_divisible(x, self.stride) for x in np.stack(shape1, 0).max(0)] # inference shape
x = [letterbox(im, new_shape=shape1 if self.pt else size, auto=False)[0] for im in imgs] # pad
x = np.stack(x, 0) if n > 1 else x[0][None] # stack
x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
shape1 = [make_divisible(x, self.stride) if self.pt else size for x in np.array(shape1).max(0)] # inf shape
x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
t.append(time_sync())

Expand Down

0 comments on commit 13051d9

Please sign in to comment.