Use tuple instead of Union in isinstance check

Author: Mainakdeb

topoloss/core.py

@@ -25,7 +25,7 @@ def __init__(
 
     def check_layer_type(self, layer):
         assert isinstance(
-            layer, Union[nn.Conv2d, nn.Linear]
+            layer, (nn.Conv2d, nn.Linear)
         ), f"Expect layer to be either nn.Conv2d or nn.Linear, but got: {type(layer)}"
 
     def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict: