updated imports of gluonts modules in pts module

awslabs · Jun 14, 2024 · de6a705 · de6a705
1 parent f0f2266
commit de6a705
Showing 1 changed file with 4 additions and 5 deletions.
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
@@ -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