watch() sometimes does not log the gradients

Describe the bug I am training a generative adversarial network. Without adversarial training, the gradients of the generator is logged. With adversarial training, the gradients of the generator is not logged.

To Reproduce Steps to reproduce the behavior:

1. Go to https://github.com/AliaksandrSiarohin/first-order-model
2. Add the following at line 71 of run.py

wandb.watch(generator, log='all')
wandb.watch(discriminator, log='all')
wandb.watch(kp_detector, log='all')

1. Run python run.py --config config/vox-256.yaml and we can see the generator and kp_detector gradients are logged.
2. Run python run.py --config config/vox-adv-256.yaml, which is adversarial training, and we can see that only discriminator gradients are logged, not those of generator and kp_detector.

Expected behavior The gradients of all the three network should be logged for adversarial training.

Screenshots

Operating System

• OS: Linux

Additional context conda list output: https://pastebin.com/sV920TQP

[vanpelt]

It’s likely regular RAM that’s the issue.

[vanpelt]

It’s likely getting killed due to memory preasure. We have to load the gradients from the GPU and if your model is really large your notebook may not have enough memory. Are you able to get a larger instance for your notebook?