fix ddp autounit UTs (pytorch#503)

Summary:
Pull Request resolved: pytorch#503

# Context
DDP UTs are broken. Reason is we're using underlying `get_swa_multi_avg_fn` which doesn't support integers (our labels) any longer. The proper solution would be to get `get_swa_multi_avg_fn` to support ints, but for now implementing this as a fix to get UTs properly running.

# This diff
Add a custom `avg_fn` to unblock tests

Reviewed By: JKSenthil

Differential Revision: D48337970

fbshipit-source-id: 760ba0fcbec746e3a4c550a5e738a7bedd782555

tests/framework/test_auto_unit.py

@@ -214,7 +214,9 @@ def forward(self, x):
                 return x
 
         my_module = Net()
-        my_swa_params = SWAParams(epoch_start=1, anneal_epochs=3)
+        my_swa_params = SWAParams(
+            epoch_start=1, anneal_epochs=3, avg_fn=lambda x, y, z: x
+        )
 
         auto_unit = DummyAutoUnit(
             module=my_module,
@@ -625,7 +627,9 @@ def forward(self, x):
                 return x
 
         my_module = Net()
-        my_swa_params = SWAParams(epoch_start=1, anneal_epochs=3)
+        my_swa_params = SWAParams(
+            epoch_start=1, anneal_epochs=3, avg_fn=lambda x, y, z: x
+        )
 
         auto_unit = DummyAutoUnit(
             module=my_module,