tests/test_sampling_utils.py

import torch

from src.model.encoder.depth.infinidepth.sampling_utils import (
    compute_image_gradient_strength,
)


def test_compute_image_gradient_strength_returns_finite_map() -> None:
    image = torch.rand(3, 8, 12)

    strength = compute_image_gradient_strength(image)

    assert strength.shape == (8, 12)
    assert torch.isfinite(strength).all()