add tests

tests/routines/test_regression.py

@@ -1,10 +1,12 @@
 # fmt:off
+from functools import partial
 from pathlib import Path
 
 from cli_test_helpers import ArgvContext
 from torch import nn
 
 from torch_uncertainty import cli_main, init_args
+from torch_uncertainty.losses import NIGLoss
 from torch_uncertainty.optimization_procedures import optim_cifar10_resnet18
 
 from .._dummies import DummyRegressionBaseline, DummyRegressionDataModule
@@ -35,6 +37,32 @@ def test_cli_main_dummy_dist(self):
 
             cli_main(model, dm, root, "dummy", args)
 
+    def test_cli_main_dummy_dist_der(self):
+        root = Path(__file__).parent.absolute().parents[0]
+        with ArgvContext("file.py"):
+            args = init_args(DummyRegressionBaseline, DummyRegressionDataModule)
+
+            # datamodule
+            args.root = str(root / "data")
+            dm = DummyRegressionDataModule(out_features=1, **vars(args))
+
+            loss = partial(
+                NIGLoss,
+                reg_weight=1e-2,
+            )
+
+            model = DummyRegressionBaseline(
+                in_features=dm.in_features,
+                out_features=4,
+                loss=loss,
+                optimization_procedure=optim_cifar10_resnet18,
+                baseline_type="single",
+                dist_estimation=4,
+                **vars(args),
+            )
+
+            cli_main(model, dm, root, "dummy_der", args)
+
     def test_cli_main_dummy(self):
         root = Path(__file__).parent.absolute().parents[0]
         with ArgvContext("file.py"):

tests/test_losses.py

@@ -4,7 +4,7 @@
 from torch import nn
 
 from torch_uncertainty.layers.bayesian import BayesLinear
-from torch_uncertainty.losses import ELBOLoss
+from torch_uncertainty.losses import ELBOLoss, NIGLoss
 
 
 # fmt: on
@@ -34,3 +34,14 @@ def test_no_bayes(self):
 
         loss = ELBOLoss(model, criterion, kl_weight=1e-5, num_samples=1)
         loss(model(torch.randn(1, 1)), torch.randn(1, 1))
+
+
+# fmt: on
+class TestNIGLoss:
+    def test_main(self):
+        with pytest.raises(ValueError):
+            NIGLoss(reg_weight=-1)
+
+        loss = NIGLoss(reg_weight=1e-2)
+
+        loss(*torch.rand((1, 4)).split(1, dim=-1), torch.randn(1, 1))

torch_uncertainty/losses.py

@@ -96,6 +96,12 @@ def forward(self, inputs: Tensor, targets: Tensor) -> Tensor:
 class NIGLoss(nn.Module):
     def __init__(self, reg_weight: float) -> None:
         super().__init__()
+
+        if reg_weight < 0:
+            raise ValueError(
+                "The regularization weight should be non-negative. "
+                f"Got {reg_weight}."
+            )
         self.reg_weight = reg_weight
 
     def _nig_nll(self, gamma, v, alpha, beta, targets):

torch_uncertainty/routines/regression.py

@@ -130,6 +130,9 @@ def validation_step(
             beta = F.softplus(logits[..., 3])
             vars = beta / (alpha - 1)
             self.val_metrics.gnll.update(means, targets, vars)
+
+            if means.ndim == 1:
+                means = means.unsqueeze(-1)
         elif self.dist_estimation == 2:
             means = logits[..., 0]
             vars = F.softplus(logits[..., 1])
@@ -162,6 +165,9 @@ def test_step(
             beta = F.softplus(logits[..., 3])
             vars = beta / (alpha - 1)
             self.test_metrics.gnll.update(means, targets, vars)
+
+            if means.ndim == 1:
+                means = means.unsqueeze(-1)
         elif self.dist_estimation == 2:
             means = logits[..., 0]
             vars = F.softplus(logits[..., 1])