Workhandle as dataclass (#1008)

facebookresearch · Jun 29, 2022 · a5116ec · a5116ec
1 parent 96fcc1c
commit a5116ec
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/fairscale/internal/params.py b/fairscale/internal/params.py
@@ -4,17 +4,18 @@
 # LICENSE file in the root directory of this source tree.
 
 import collections.abc as abc
+from dataclasses import dataclass
 from math import inf
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
 import torch.distributed as dist
 
 
+@dataclass
 class Workhandle:
-    def __init__(self, handle: Any, callback: Optional[Callable]) -> None:
-        self.handle = handle
-        self.callback = callback
+    handle: Any
+    callback: Optional[Callable] = None
 
 
 def get_global_rank(group: Any, rank: int) -> int: