Hey Simon, so I rewrote my custom layer. But I still seem to be “corrupting” the gradient somewhere. Is there a particular way you would suggest looking for this mistake? Something like the following code. I’m just trying to track it down now 

    def forward(self, x):
        # Regular layer
        x = F.relu(F.max_pool2d(self.conv1(x), 2))

        # Me trying to get a gradient
        z = torch.add(x, 1)
        s = torch.mul(z, z)
        out = s.mean()
        out.backward()
        print(x.grad)