chore: update return types for train methods

Author: abelaba

sbi/inference/trainers/base.py

@@ -21,7 +21,7 @@
 from warnings import warn
 
 import torch
-from torch import Tensor, nn
+from torch import Tensor
 from torch.distributions import Distribution
 from torch.nn.utils.clip_grad import clip_grad_norm_
 from torch.optim.adam import Adam
@@ -311,7 +311,7 @@ def train(
         discard_prior_samples: bool = False,
         retrain_from_scratch: bool = False,
         show_train_summary: bool = False,
-    ) -> NeuralPosterior: ...
+    ) -> Union[ConditionalEstimator, RatioEstimator]: ...
 
     @abstractmethod
     def _initialize_neural_network(
@@ -906,7 +906,7 @@ def _run_training_loop(
         show_train_summary: bool,
         loss_kwargs: Optional[Dict[str, Any]] = None,
         summarization_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> nn.Module:
+    ) -> Union[ConditionalEstimator, RatioEstimator]:
         """
         Run the main training loop for the neural network, including epoch-wise
         training, validation, and convergence checking.
@@ -940,6 +940,9 @@ def _run_training_loop(
 
         assert self._neural_net is not None
 
+        # Move entire net to device for training.
+        self._neural_net.to(self._device)
+
         if not resume_training:
             self.optimizer = Adam(
                 list(self._neural_net.parameters()),

sbi/inference/trainers/npe/npe_base.py

@@ -319,7 +319,6 @@ def default_calibration_kernel(x):
             calibration_kernel=calibration_kernel,
             force_first_round_loss=force_first_round_loss,
         )
-
         return self._run_training_loop(  # type: ignore
             train_loader=train_loader,
             val_loader=val_loader,
@@ -645,9 +644,6 @@ def _initialize_neural_network(
 
             del theta, x
 
-        # Move entire net to device for training.
-        self._neural_net.to(self._device)
-
     def _get_training_losses(self, batch: Any, loss_kwargs: Dict[str, Any]) -> Tensor:
         """
         Compute training losses for a batch of data.

sbi/inference/trainers/nre/nre_base.py

@@ -6,7 +6,7 @@
 from typing import Any, Dict, Literal, Optional, Protocol, Tuple, Union
 
 import torch
-from torch import Tensor, eye, nn, ones
+from torch import Tensor, eye, ones
 from torch.distributions import Distribution
 from torch.utils.tensorboard.writer import SummaryWriter
 from typing_extensions import Self
@@ -175,7 +175,7 @@ def train(
         show_train_summary: bool = False,
         dataloader_kwargs: Optional[Dict] = None,
         loss_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> nn.Module:
+    ) -> RatioEstimator:
         r"""Return classifier that approximates the ratio $p(\theta,x)/p(\theta)p(x)$.
 
         Args:
@@ -237,7 +237,7 @@ def train(
 
         loss_kwargs["num_atoms"] = num_atoms
 
-        return self._run_training_loop(
+        return self._run_training_loop(  # type: ignore
             train_loader=train_loader,
             val_loader=val_loader,
             max_num_epochs=max_num_epochs,
@@ -428,8 +428,6 @@ def _initialize_neural_network(
 
             del x, theta
 
-        self._neural_net.to(self._device)
-
     def _get_training_losses(self, batch: Any, loss_kwargs: Dict[str, Any]) -> Tensor:
         """
         Compute training losses for a batch of data.

sbi/inference/trainers/vfpe/base_vf_inference.py

@@ -667,6 +667,3 @@ def _initialize_neural_network(
             )
 
             del theta, x
-
-        # Move entire net to device for training.
-        self._neural_net.to(self._device)