core: pass strict_layer_type arg to cortical sheet op

Author: Mayukhdeb

topoloss/core.py

@@ -39,9 +39,9 @@ def get_layerwise_topo_losses(self, model, do_scaling: bool = True) -> dict:
 
             if isinstance(loss_info, LaplacianPyramid):
                 if isinstance(layer, nn.Linear):
-                    cortical_sheet = get_cortical_sheet_linear(layer=layer)
+                    cortical_sheet = get_cortical_sheet_linear(layer=layer, strict_layer_type=self.strict_layer_type)
                 else:
-                    cortical_sheet = get_cortical_sheet_conv(layer=layer)
+                    cortical_sheet = get_cortical_sheet_conv(layer=layer, strict_layer_type=self.strict_layer_type)
 
                 loss = laplacian_pyramid_loss(
                     cortical_sheet=cortical_sheet,

topoloss/cortical_sheet/input.py

@@ -3,8 +3,9 @@
 from .common import find_cortical_sheet_size, GridDimensions2D
 
 
-def get_cortical_sheet_linear_input(layer: nn.Linear):
-    assert isinstance(layer, nn.Linear)
+def get_cortical_sheet_linear_input(layer: nn.Linear, strict_layer_type: bool = True):
+    if strict_layer_type is True:
+        assert isinstance(layer, nn.Linear)
     weight = layer.weight
     num_input_neurons = weight.shape[1]
     assert weight.ndim == 2

topoloss/cortical_sheet/output.py

@@ -3,8 +3,9 @@
 from .common import find_cortical_sheet_size, GridDimensions2D
 
 
-def get_cortical_sheet_linear(layer: nn.Linear):
-    assert isinstance(layer, nn.Linear)
+def get_cortical_sheet_linear(layer: nn.Linear, strict_layer_type: bool):
+    if strict_layer_type is True:
+        assert isinstance(layer, nn.Linear)
     weight = layer.weight
     num_output_neurons = weight.shape[0]
     assert weight.ndim == 2
@@ -16,8 +17,9 @@ def get_cortical_sheet_linear(layer: nn.Linear):
     )
 
 
-def get_cortical_sheet_conv(layer: nn.Conv2d):
-    assert isinstance(layer, nn.Conv2d)
+def get_cortical_sheet_conv(layer: nn.Conv2d, strict_layer_type: bool):
+    if strict_layer_type is True:
+        assert isinstance(layer, nn.Conv2d)
     weight = layer.weight
     assert weight.ndim == 4
     num_output_channels = weight.shape[0]