Skip to content

GitLab

You need to sign in or sign up before continuing.
Unverified Commit a5116ecd authored by Crutcher Dunnavant's avatar Crutcher Dunnavant Committed by GitHub
Browse files

Workhandle as dataclass (#1008)

parent 96fcc1ca
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 3 deletions
fairscale/internal/params.py
View file @ a5116ecd
...@@ -4,6 +4,7 @@ ...@@ -4,6 +4,7 @@
# LICENSE file in the root directory of this source tree. # LICENSE file in the root directory of this source tree.
import collections.abc as abc import collections.abc as abc
from dataclasses import dataclass
from math import inf from math import inf
from typing import Any, Callable, Dict, List, Optional from typing import Any, Callable, Dict, List, Optional
...@@ -11,10 +12,10 @@ import torch ...@@ -11,10 +12,10 @@ import torch
import torch.distributed as dist import torch.distributed as dist
@dataclass
class Workhandle: class Workhandle:
def __init__(self, handle: Any, callback: Optional[Callable]) -> None: handle: Any
self.handle = handle callback: Optional[Callable] = None
self.callback = callback
def get_global_rank(group: Any, rank: int) -> int: def get_global_rank(group: Any, rank: int) -> int:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment