Squeezenet reshape outputs fix (#10222)

@AyushExel Signed-off-by: Glenn Jocher <[email protected]> Signed-off-by: Glenn Jocher <[email protected]>
ultralytics · Nov 19, 2022 · 72cad39 · 72cad39
1 parent 40bb803
commit 72cad39
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -82,7 +82,7 @@ def reshape_classifier_output(model, n=1000):
         elif nn.Conv2d in types:
             i = types.index(nn.Conv2d)  # nn.Conv2d index
             if m[i].out_channels != n:
-                m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias)
+                m[i] = nn.Conv2d(m[i].in_channels, n, m[i].kernel_size, m[i].stride, bias=m[i].bias is not None)
 
 
 @contextmanager