Hi,
You are running a ReLU after computing final output. If the initial output is all negative, this will return a zero-tensor & zero gradients.
Change
return F.relu(self.output_layer(x)) to return self.output_layer(x)