(
self,
model: torch.nn.Module,
train_data: DataLoader,
optimizer: torch.optim.Optimizer,
save_every: int,
snapshot_path: str,
)
| 16 | |
| 17 | class Trainer: |
| 18 | def __init__( |
| 19 | self, |
| 20 | model: torch.nn.Module, |
| 21 | train_data: DataLoader, |
| 22 | optimizer: torch.optim.Optimizer, |
| 23 | save_every: int, |
| 24 | snapshot_path: str, |
| 25 | ) -> None: |
| 26 | self.gpu_id = int(os.environ["LOCAL_RANK"]) |
| 27 | self.model = model.to(self.gpu_id) |
| 28 | self.train_data = train_data |
| 29 | self.optimizer = optimizer |
| 30 | self.save_every = save_every |
| 31 | self.epochs_run = 0 |
| 32 | self.snapshot_path = snapshot_path |
| 33 | if os.path.exists(snapshot_path): |
| 34 | print("Loading snapshot") |
| 35 | self._load_snapshot(snapshot_path) |
| 36 | |
| 37 | self.model = DDP(self.model, device_ids=[self.gpu_id]) |
| 38 | |
| 39 | def _load_snapshot(self, snapshot_path): |
| 40 | loc = f"cuda:{self.gpu_id}" |
nothing calls this directly
no test coverage detected