Allow variable number of input channels

models/spatial_model.py

@@ -3,10 +3,10 @@ from .blocks import SpatialResBlock
 
 
 class SpatialResNet(nn.Module):
-    def __init__(self, n_classes):
+    def __init__(self, n_channels, n_classes):
         super(SpatialResNet, self).__init__()
 
-        self.block1 = SpatialResBlock(in_channels=1, out_channels=16, downsample=False)
+        self.block1 = SpatialResBlock(in_channels=n_channels, out_channels=16, downsample=False)
         self.block2 = SpatialResBlock(in_channels=16, out_channels=32)
         self.block3 = SpatialResBlock(in_channels=32, out_channels=64)