stubs/torch/optim/__init__.py

"""Minimal stub for torch.optim used by the project."""
from typing import Iterable

class Optimizer:
    def __init__(self, params: Iterable, lr: float = 0.001):
        self.params = params
        self.lr = lr

class Adam(Optimizer):
    def __init__(self, params: Iterable, lr: float = 0.001):
        super().__init__(params, lr)

def AdamW(params: Iterable, lr: float = 0.001):
    return Adam(params, lr)