Verified Commit b59d8c91 authored by Max Ehrlich's avatar Max Ehrlich
Browse files

Allow variable number of input channels

parent 20c01fcc
......@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment