Solved by ptrblck in post #2 
              
              
                Since you are using single channel images, you could swap the batch and the channel dimension and use a grouped convolution: 
class FilterNet(nn.Module):

    def __init__(self, obs_size, hidden_size, kernel_size, activation_func='tanh'):
        super(FilterNet, self).__init__()
        self.kernel…