import argparse

import torch
from torch import nn


def check_finite(name: str, tensor: torch.Tensor, *, step: int) -> None:
    finite = torch.isfinite(tensor)
    if bool(finite.all()):
        return

    values = tensor.detach()[~finite].flatten()[:4].tolist()
    print(f"non_finite={name}")
    print(f"step={step} shape={tuple(tensor.shape)} dtype={tensor.dtype}")
    print(f"values={values}")
    raise SystemExit(1)