diff --git a/src/gluonts/nursery/robust-mts-attack/pts/modules/distribution_output.py b/src/gluonts/nursery/robust-mts-attack/pts/modules/distribution_output.py
index 9b6a4cb157..dd528262af 100644
--- a/src/gluonts/nursery/robust-mts-attack/pts/modules/distribution_output.py
+++ b/src/gluonts/nursery/robust-mts-attack/pts/modules/distribution_output.py
@@ -44,11 +44,10 @@
     TransformedImplicitQuantile,
 )
 from gluonts.core.component import validated
-from gluonts.torch.modules.distribution_output import (
-    DistributionOutput,
-    LambdaLayer,
-    PtArgProj,
-)
+from gluonts.torch.distributions.distribution_output import DistributionOutput
+from gluonts.torch.modules.lambda_layer import LambdaLayer
+from gluonts.torch.distributions.output import PtArgProj
+
 from pts.modules.iqn_modules import ImplicitQuantileModule