from typing import Optional
import torch


def check_nans_(x, warn: Optional[str] = None, value: float = 0):
    """Mask out all non-finite values + warn if `warn is not None`"""
    msk = torch.isfinite(x)
    if warn is not None:
        if ~(msk.all()):
            print(f'WARNING: NaNs in {warn}')
    x.masked_fill_(msk.bitwise_not(), value)
    return x