def test_ssim_loss_grad(x_y_4d_5d, device: str) -> None:
x = x_y_4d_5d[0].to(device)
y = x_y_4d_5d[1].to(device)
x.requires_grad_(True)
loss = SSIMLoss(data_range=1.)(x, y).mean()
loss.backward()
assert torch.isfinite(
x.grad).all(), f'Expected finite gradient values, got {x.grad}'
def test_ssim_loss_grad(prediction_target_4d_5d: Tuple[torch.Tensor,
torch.Tensor],
device: str) -> None:
prediction = prediction_target_4d_5d[0].to(device)
target = prediction_target_4d_5d[1].to(device)
prediction.requires_grad_(True)
loss = SSIMLoss(data_range=1.)(prediction, target).mean()
loss.backward()
assert torch.isfinite(prediction.grad).all(
), f'Expected finite gradient values, got {prediction.grad}'
