fix imports.

sbi-dev · Feb 9, 2024 · 32c6b3a · 32c6b3a
1 parent ce04fa0
commit 32c6b3a
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 4 deletions.
diff --git a/sbi/inference/snle/snle_base.py b/sbi/inference/snle/snle_base.py
@@ -11,7 +11,6 @@
 from torch import Tensor, nn, optim
 from torch.distributions import Distribution
 from torch.nn.utils.clip_grad import clip_grad_norm_
-from torch.utils import data
 from torch.utils.tensorboard.writer import SummaryWriter
 
 from sbi import utils as utils

diff --git a/sbi/inference/snpe/snpe_base.py b/sbi/inference/snpe/snpe_base.py
@@ -26,14 +26,15 @@
     RestrictedPrior,
     check_estimator_arg,
     handle_invalid_x,
+    mask_sims_from_prior,
     nle_nre_apt_msg_on_invalid_x,
     npe_msg_on_invalid_x,
     test_posterior_net_for_multi_d_x,
     validate_theta_and_x,
     warn_if_zscoring_changes_data,
     x_shape_from_simulation,
 )
-from sbi.utils.sbiutils import ImproperEmpirical, mask_sims_from_prior
+from sbi.utils.sbiutils import ImproperEmpirical
 
 
 class PosteriorEstimator(NeuralInference, ABC):

diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
@@ -12,9 +12,8 @@
 import torch
 import torch.distributions.transforms as torch_tf
 from pyro.distributions import Empirical
-from torch import Tensor
+from torch import Tensor, ones, optim, zeros
 from torch import nn as nn
-from torch import ones, optim, zeros
 from torch.distributions import Distribution, Independent, biject_to, constraints
 
 from sbi import utils as utils