Increase FLOPS robustness (#1589)

2020-12-05 11:41:17 +01:00 · 2020-12-05 11:41:17 +01:00 · dbcb192f2d
commit dbcb192f2d
parent d1ad63206b
1 changed files with 1 additions and 1 deletions
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@ -152,7 +152,7 @@ def model_info(model, verbose=False, img_size=640):

    try:  # FLOPS
        from thop import profile
-        stride = int(model.stride.max())
+        stride = int(model.stride.max()) if hasattr(model, 'stride') else 32
        img = torch.zeros((1, 3, stride, stride), device=next(model.parameters()).device)  # input
        flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2  # stride FLOPS
        img_size = img_size if isinstance(img_size, list) else [img_size, img_size]  # expand if int/float