Saving custom layer parameters

2

You can use the track_grad_norm flag of the trainer to ensure that the custom layer is receiving gradients.

1 Like