def forward(self, input_tensor):
# Swap axis of RGB image for the network to get a "batch" of size = 3 rather the 3 channels
input_tensor = swap_axis(input_tensor)
downscaled = self.first_layer(input_tensor)
features = self.feature_block(downscaled)
output = self.final_layer(features)
return swap_axis(output)
def forward(self, input_tensor):
input_tensor = swap_axis(input_tensor)
downscaled = self.first_layer(input_tensor)
features = self.feature_block(downscaled)
output = self.final_layer(features) + self.downsample(input_tensor)
return swap_axis(output)