import torch


def normalise_grad(grad: torch.Tensor) -> torch.Tensor:
    """Normalise and handle NaNs caused by norm division.

    Args:
        grad (torch.Tensor): Gradient to normalise

    Returns:
        torch.Tensor: Normalised gradient
    """
    normed_grad = grad / torch.linalg.norm(grad, dim=1)[:, None]
    if torch.isnan(normed_grad).any():
        normed_grad = torch.where(
            torch.isnan(normed_grad),
            torch.zeros_like(normed_grad),
            normed_grad,
        )
    return normed_grad