diff --git a/switch/estimator.py b/switch/estimator.py
index bd08c16..bfe9dd5 100644
--- a/switch/estimator.py
+++ b/switch/estimator.py
@@ -8,7 +8,7 @@ from gluonts.itertools import Cyclic, IterableSlice, PseudoShuffled
 from gluonts.time_feature import TimeFeature, time_features_from_frequency_str
 from gluonts.torch.model.estimator import PyTorchLightningEstimator
 from gluonts.torch.model.predictor import PyTorchPredictor
-from gluonts.torch.modules.distribution_output import DistributionOutput, StudentTOutput
+from gluonts.torch.distributions import DistributionOutput, StudentTOutput
 from gluonts.torch.modules.loss import DistributionLoss, NegativeLogLikelihood
 from gluonts.torch.util import IterableDataset
 from gluonts.transform import (