18 lines
420 B
Python
18 lines
420 B
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class Job:
|
||
|
|
script: str
|
||
|
|
num_nodes: int = 2
|
||
|
|
min_vram: int = 8
|
||
|
|
docker_image: str = "pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime"
|
||
|
|
rounds: int = 10
|
||
|
|
local_steps: int = 100
|
||
|
|
lr: float = 0.01
|
||
|
|
batch_size: int = 64
|
||
|
|
checkpoint_dir: str | None = None
|
||
|
|
env: dict[str, str] = field(default_factory=dict)
|