[Fix] - Finiteness check for all tensors (#890)

* Finiteness check for all tensors

* [skip ci] updating changelog

CHANGELOG.md

@@ -5,6 +5,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [0.4.4] - TBD
+
+### Fixed
+- inf/nan check on all gradient tensors in ShardedGradScaler [#890]
+
 ### Added
 
 ### Changed

fairscale/optim/grad_scaler.py

@@ -156,7 +156,7 @@ class ShardedGradScaler(TorchGradScaler):
         assert found_inf.numel() == 1, "found_inf must be a 1-element tensor."
 
         expected_device = grads[0].device
-        for tensor in grads[0]:
+        for tensor in grads:
             try:
                 assert tensor.device == expected_device, "grads must be on the same device"
             except AssertionError: