small fixes to layerwise gradient scaler (#910)

fairscale/optim/layerwise_gradient_scaler.py

@@ -200,7 +200,7 @@ class LayerwiseGradientScaler:
         layers_with_finite_values = self._get_layers_with_finite_values()
         for item in layers_with_finite_values:
             for param_name, param in item.layer.named_parameters():
-                if hasattr(param, "grad"):
+                if hasattr(param, "grad") and param.grad is not None:
                     logging.debug("%s scaling down %s by %s" % (item.layer_name, param_name, 1.0 / item.scaling_factor))
                     param.grad.mul_(1.0 / item.scaling_factor)
 

tests/optim/test_layerwise_gradient_scaler.py

 import logging
+import os
 from typing import Any, List, Tuple, Union
 
 import numpy as np
@@ -200,6 +201,7 @@ def test_vision_model() -> None:
     # Remove randomness from various sources while testing.
     torch.use_deterministic_algorithms(True)  # type: ignore
     # set environment variable in CircleCI for test to pass: CUBLAS_WORKSPACE_CONFIG = :4096:8
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
 
     m1 = SimpleConvNet()
     m2 = SimpleConvNet()