LARS exclude BN: use a flag instead of a function (#388)

Summary:
This fix the problem that the Barlow Twins model needs to save a function in the checkpoint.

Pull Request resolved: #388

Reviewed By: iseessel

Differential Revision: D30158877

Pulled By: prigoyal

fbshipit-source-id: 537d0686422148447a4a42e14b448eb6e592eec9

Author: jingli9111

vissl/optimizers/lars.py

@@ -87,7 +87,7 @@ def __init__(
             "weight_decay": weight_decay,
             "momentum": momentum,
             "eta": eta,
-            "exclude": self._exclude_bias_and_norm if exclude_bias_and_norm else None,
+            "exclude": exclude_bias_and_norm,
         }
         super().__init__(params, defaults)
 
@@ -106,7 +106,7 @@ def step(self):
 
                 dp = dp.add(p, alpha=g["weight_decay"])
 
-                if g["exclude"] is None or not g["exclude"](p):
+                if not g["exclude"] or not self._exclude_bias_and_norm(p):
                     param_norm = torch.norm(p)
                     update_norm = torch.norm(dp)