mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
b2953f5643
This is follow-up of #165214 to continue applying ruff UP035 rule to the code base. Pull Request resolved: https://github.com/pytorch/pytorch/pull/165515 Approved by: https://github.com/Lucaskabela
28 lines
544 B
Python
28 lines
544 B
Python
import dataclasses
|
|
from collections.abc import Callable
|
|
from typing import Optional
|
|
|
|
|
|
all_experiments: dict[str, Callable] = {}
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Experiment:
|
|
name: str
|
|
metric: str
|
|
target: float
|
|
actual: float
|
|
dtype: str
|
|
device: str
|
|
arch: str # GPU name for CUDA or CPU arch for CPU
|
|
is_model: bool = False
|
|
|
|
|
|
def register_experiment(name: Optional[str] = None):
|
|
def decorator(func):
|
|
key = name or func.__name__
|
|
all_experiments[key] = func
|
|
return func
|
|
|
|
return decorator
|