glenn-jocher commited on
Commit
84efa62
1 Parent(s): caf7ad0

Fix PyTorch Hub export inference shapes (#6949)

Browse files

May resolve https://github.com/ultralytics/yolov5/issues/6947

Files changed (1) hide show
  1. models/common.py +3 -4
models/common.py CHANGED
@@ -544,10 +544,9 @@ class AutoShape(nn.Module):
544
  g = (size / max(s)) # gain
545
  shape1.append([y * g for y in s])
546
  imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
547
- shape1 = [make_divisible(x, self.stride) for x in np.stack(shape1, 0).max(0)] # inference shape
548
- x = [letterbox(im, new_shape=shape1 if self.pt else size, auto=False)[0] for im in imgs] # pad
549
- x = np.stack(x, 0) if n > 1 else x[0][None] # stack
550
- x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
551
  x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
552
  t.append(time_sync())
553
 
 
544
  g = (size / max(s)) # gain
545
  shape1.append([y * g for y in s])
546
  imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update
547
+ shape1 = [make_divisible(x, self.stride) if self.pt else size for x in np.array(shape1).max(0)] # inf shape
548
+ x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
549
+ x = np.ascontiguousarray(np.array(x).transpose((0, 3, 1, 2))) # stack and BHWC to BCHW
 
550
  x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32
551
  t.append(time_sync())
552