Changed LazyAwaitable.__torch_function__ to clasmethod (pytorch#2270)

Summary:
Pull Request resolved: pytorch#2270

This avoids the deprecation warning causing a lot of log spam.

Reviewed By: idning, yjhao, sarckk, wanchaol

Differential Revision: D60775821

fbshipit-source-id: d42ad58578cf3756e8d854db70b2877263d07602

torchrec/distributed/types.py

@@ -371,8 +371,9 @@ def _wait_async(obj: Any) -> Any:
         else:
             return obj
 
+    @classmethod
     # pyre-ignore [2, 3]
-    def __torch_function__(self, func, types, args=(), kwargs=None):
+    def __torch_function__(cls, func, types, args=(), kwargs=None):
         """
         The LazyAwaitable type has a `__torch_function__` implementation.
         This means when this type is seens as an argument to a PyTorch