New smart_inference_mode() conditional decorator (#8957)

New smart_inference_mode()
This commit is contained in:
Glenn Jocher authored and GitHub committed 2022-08-13 20:38:51 +02:00
1 parent 6aed0a7c00
commit dc38cd03f4
6 files changed
+26 -18

No files matched your search

+2 -2
View File
@@ -76,12 +76,12 @@ class Detect(nn.Module):
return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
def _make_grid(self, nx=20, ny=20, i=0):
def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, '1.10.0')):
d = self.anchors[i].device
t = self.anchors[i].dtype
shape = 1, self.na, ny, nx, 2 # grid shape
y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t)
if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
if torch_1_10: # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility
yv, xv = torch.meshgrid(y, x, indexing='ij')
else:
yv, xv = torch.meshgrid(y, x)