new_grad_tensor.copy_(param.grad.view(-1)) AttributeError: 'NoneType' object has no attribute 'view'

I'm trying to apply deepspeed stage 2 to stylegan2 but I get this error.

Here's my config:

{
  "train_batch_size" : 4,
  "optimizer": {
    "type": "Adam",
    "params": {
      "lr": 0.0002,
      "betas": [
        0.5,
        0.999
      ],
      "eps": 1e-8
    }
  },
  "steps_per_print" : 10,
  "fp16": {
      "enabled": true
  },
  "zero_optimization": {
      "stage": 2,
      "cpu_offload": true,
      "contiguous_gradients": true,
      "overlap_comm": true
  }
}

And here's the full stack trace:

Traceback (most recent call last):
  File "stylegan2_pytorch/ucl_deepspeed.py", line 200, in <module>
    main()
  File "stylegan2_pytorch/ucl_deepspeed.py", line 197, in main
    train_from_folder(deepspeed_args=deepspeed_args)
  File "stylegan2_pytorch/ucl_deepspeed.py", line 177, in train_from_folder
    run_training(0, 1, model_args, data, load_from, new, num_train_steps, name, seed)
  File "stylegan2_pytorch/ucl_deepspeed.py", line 62, in run_training
    retry_call(model.train, tries=3, exceptions=NanException)
  File "/opt/conda/lib/python3.7/site-packages/retry/api.py", line 101, in retry_call
    return __retry_internal(partial(f, *args, **kwargs), exceptions, tries, delay, max_delay, backoff, jitter, logger)
  File "/opt/conda/lib/python3.7/site-packages/retry/api.py", line 33, in __retry_internal
    return f()
  File "/home/dtkatch/stylegan2-pytorch/stylegan2_pytorch/stylegan2_pytorch.py", line 1052, in train
    self.GAN.model_engineG.backward(gen_loss)
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/engine.py", line 845, in backward
    self.optimizer.backward(loss)
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 1609, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 53, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/opt/conda/lib/python3.7/site-packages/torch/tensor.py", line 221, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/opt/conda/lib/python3.7/site-packages/torch/autograd/__init__.py", line 132, in backward
    allow_unreachable=True)  # allow_unreachable flag
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 594, in reduce_partition_and_remove_grads
    self.reduce_ready_partitions_and_remove_grads(param, i)
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 984, in reduce_ready_partitions_and_remove_grads
    self.reduce_independent_p_g_buckets_and_remove_grads(param, i)
  File "/opt/conda/lib/python3.7/site-packages/deepspeed/runtime/zero/stage2.py", line 637, in reduce_independent_p_g_buckets_and_remove_grads
    new_grad_tensor.copy_(param.grad.view(-1))
AttributeError: 'NoneType' object has no attribute 'view'

Start with deepspeed/runtime/zero/stage2.py, especially reduce_independent_p_g_buckets_and_remove_grads and the call path shown in the traceback. Reproduce using the provided ZeRO stage 2 configuration with stylegan2, determine why param.grad is None, and confirm training completes without the AttributeError.