core: assert isinstance linear if strict_layer_type is True

Author: Mayukhdeb

topoloss/core.py

@@ -69,9 +69,10 @@ def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict:
                 )
 
             elif isinstance(loss_info, LaplacianPyramidOnInput):
-                assert isinstance(
-                    layer, nn.Linear
-                ), f"Expected layer to be nn.Linear, but got: {type(layer)}"
+                if self.strict_layer_type is True:
+                    assert isinstance(
+                        layer, nn.Linear
+                    ), f"Expected layer to be nn.Linear, but got: {type(layer)}"
                 cortical_sheet = get_cortical_sheet_linear_input(layer=layer)
                 loss = laplacian_pyramid_loss(
                     cortical_sheet=cortical_sheet,