I think I squashed the bugs?

retinal_rl/analysis/statistics.py

@@ -192,6 +192,7 @@ def cnn_statistics(
     # Analyze input statistics
     input_spectral = _layer_spectral_analysis(device, dataloader, nn.Identity())
     input_histograms = _layer_pixel_histograms(device, dataloader, nn.Identity())
+    # num channels as FloatArray
 
     # Load input statistics into results dictionary
     results["input"] = {

retinal_rl/models/circuits/convolutional.py

@@ -47,6 +47,7 @@ def __init__(
         stride: Union[int, List[int]],
         activation: Union[str, List[str]],
         normalization: bool = False,
+        affine_norm: bool = True,
         layer_names: Optional[List[str]] = None,
     ):
         super().__init__(input_shape)
@@ -85,7 +86,12 @@ def __init__(
             )
             if normalization and i is num_layers - 1:
                 conv_layers.append(
-                    ("instance_norm", nn.InstanceNorm2d(self.num_channels[i]))
+                    (
+                        f"{layer_names[i]}_instance_norm"
+                        if layer_names is not None
+                        else "instance_norm" + str(i),
+                        nn.InstanceNorm2d(self.num_channels[i], affine=affine_norm),
+                    )
                 )
             conv_layers.append((actnm, self.str_to_activation(self.activation[i])))
         self.conv_head = nn.Sequential(OrderedDict(conv_layers))
@@ -106,6 +112,7 @@ def __init__(
         stride: Union[int, List[int]],
         activation: Union[str, List[str]],
         normalization: bool = False,
+        affine_norm: bool = True,
         layer_names: Optional[List[str]] = None,
     ):
         # add parameters to model and apply changes for internal use
@@ -153,7 +160,12 @@ def __init__(
             )
             if normalization and i == 0:
                 deconv_layers.append(
-                    ("instance_norm", nn.InstanceNorm2d(self.num_channels[i]))
+                    (
+                        f"{layer_names[i]}_instance_norm"
+                        if layer_names is not None
+                        else "instance_norm" + str(i),
+                        nn.InstanceNorm2d(self.num_channels[i], affine=affine_norm),
+                    )
                 )
 
             deconv_layers.append((actnm, self.str_to_activation(self.activation[i])))

runner/util.py

@@ -120,7 +120,6 @@ def assemble_neural_circuits(
         if node in sensor_shapes:
             continue
 
-        print(f"Processing node {node}")
         circuit_config = OmegaConf.select(circuits, node)
         input_tensor = _assemble_inputs(node, connectome, dummy_responses)
         input_shape = list(input_tensor.shape[1:])