2cc34fb72b
Closes #914 Additional changes: - Deprecate python below 11 - Remove 3rd party and throughput tests. This simplifies install and test pipeline - Remove gym compatibility and shimmy - Format with 3.11 conventions. In particular, add `zip(..., strict=True/False)` where possible Since the additional tests and gym were complicating the CI pipeline (flaky and dist-dependent), it didn't make sense to work on fixing the current tests in this PR to then just delete them in the next one. So this PR changes the build and removes these tests at the same time.
23 lines
517 B
Python
23 lines
517 B
Python
import torch
|
|
from torch import nn
|
|
|
|
|
|
def optim_step(
|
|
loss: torch.Tensor,
|
|
optim: torch.optim.Optimizer,
|
|
module: nn.Module,
|
|
max_grad_norm: float | None = None,
|
|
) -> None:
|
|
"""Perform a single optimization step.
|
|
|
|
:param loss:
|
|
:param optim:
|
|
:param module:
|
|
:param max_grad_norm: if passed, will clip gradients using this
|
|
"""
|
|
optim.zero_grad()
|
|
loss.backward()
|
|
if max_grad_norm:
|
|
nn.utils.clip_grad_norm_(module.parameters(), max_norm=max_grad_norm)
|
|
optim.step()
|