🔨 better Exception

Lightning-AI · Dec 1, 2020 · dfece30 · dfece30
1 parent d2f0a02
commit dfece30
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pytorch_lightning/core/memory.py b/pytorch_lightning/core/memory.py
@@ -224,8 +224,8 @@ def _get_out_size_params(out_sizes):
                 try:
                     # try to find prod, i.e check for unknown sizes.
                     _total_out_params += np.prod(out_sizes)
-                except:
-                    # do nothing if could not find product.
+                except TypeError:
+                    # do nothing if tried to find prod on unknown type.
                     pass
             else:
                 _ = [_get_out_size_params(out_size) for out_size in out_sizes if isinstance(out_size, list)]