# Buggy code
def create_transformation(self) -> Transformation:
    remove_field_names = [
        FieldName.FEAT_DYNAMIC_CAT,
        FieldName.FEAT_STATIC_REAL,
    ]
    if not self.use_feat_dynamic_real:
        remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)

    return Chain(
        [RemoveFields(field_names=remove_field_names)]
        + (
            [SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0.0])]
            if not self.use_feat_static_cat
            else []
        )
        + [
            AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1),
            AsNumpyArray(field=FieldName.TARGET, expected_ndim=1),
            # gives target the (1, T) layout
            ExpandDimArray(field=FieldName.TARGET, axis=0),
            AddObservedValuesIndicator(
                target_field=FieldName.TARGET,
                output_field=FieldName.OBSERVED_VALUES,
            ),
            # Unnormalized seasonal features
            AddTimeFeatures(
                time_features=CompositeISSM.seasonal_features(self.freq),
                pred_length=self.prediction_length,
                start_field=FieldName.START,
                target_field=FieldName.TARGET,
                output_field=SEASON_INDICATORS_FIELD,
            ),
            AddTimeFeatures(
                start_field=FieldName.START,
                target_field=FieldName.TARGET,
                output_field=FieldName.FEAT_TIME,
                time_features=self.time_features,
                pred_length=self.prediction_length,
            ),
            AddAgeFeature(
                target_field=FieldName.TARGET,
                output_field=FieldName.FEAT_AGE,
                pred_length=self.prediction_length,
                log_scale=True,
            ),
            VstackFeatures(
                output_field=FieldName.FEAT_TIME,
                input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
                + ([FieldName.FEAT_DYNAMIC_REAL] if self.use_feat_dynamic_real else []),
            ),
            CanonicalInstanceSplitter(
                target_field=FieldName.TARGET,
                is_pad_field=FieldName.IS_PAD,
                start_field=FieldName.START,
                forecast_start_field=FieldName.FORECAST_START,
                instance_sampler=TestSplitSampler(),
                time_series_fields=[
                    FieldName.FEAT_TIME,
                    SEASON_INDICATORS_FIELD,
                    FieldName.OBSERVED_VALUES,
                ],
                allow_target_padding=True,
                instance_length=self.past_length,
                use_prediction_features=True,
                prediction_length=self.prediction_length,
            ),
        ]
    )


# Traceback (mxnet.base.MXNetError)
Traceback (most recent call last):
File "D:/python/DeepSSM/deep ISSM.py", line 53, in <module>
predictor = estimator.train(dataset.train)
File "D:\python\DeepSSM\gluonts\model\estimator.py", line 254, in train
training_data, validation_data, num_workers, num_prefetch, **kwargs
File "D:\python\DeepSSM\gluonts\model\estimator.py", line 233, in train_model
validation_iter=validation_data_loader,
File "D:\python\DeepSSM\gluonts\trainer\_base.py", line 308, in __call__
epoch_loss = loop(epoch_no, train_iter)
File "D:\python\DeepSSM\gluonts\trainer\_base.py", line 242, in loop
output = net(*inputs)
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\gluon\block.py", line 540, in __call__
out = self.forward(*args)
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\gluon\block.py", line 917, in forward
return self.hybrid_forward(ndarray, x, *args, **params)
File "D:\python\DeepSSM\gluonts\model\deepstate\_network.py", line 193, in hybrid_forward
length=self.past_length,
File "D:\python\DeepSSM\gluonts\model\deepstate\_network.py", line 151, in compute_lds
) = self.issm.get_issm_coeff(seasonal_indicators)
File "D:\python\DeepSSM\gluonts\model\deepstate\issm.py", line 132, in get_issm_coeff
self.innovation_coeff(seasonal_indicators),
File "D:\python\DeepSSM\gluonts\model\deepstate\issm.py", line 247, in innovation_coeff
axis=2
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\ndarray\ndarray.py", line 1695, in squeeze
return op.squeeze(self, *args, **kwargs)
File "<string>", line 45, in squeeze
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\_ctypes\ndarray.py", line 92, in _imperative_invoke
ctypes.byref(out_stypes)))
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\base.py", line 252, in check_call
raise MXNetError(py_str(_LIB.MXGetLastError()))
mxnet.base.MXNetError: [13:47:56] c:\jenkins\workspace\mxnet-tag\mxnet\src\operator\tensor\./matrix_op-inl.h:2178: Check failed: dshape[axes[i]] == 1 (2 vs. 1) cannot select an axis to squeeze out which has size=2 not equal to one

# Fixed code
def create_transformation(self) -> Transformation:
    remove_field_names = [
        FieldName.FEAT_DYNAMIC_CAT,
        FieldName.FEAT_STATIC_REAL,
    ]
    if not self.use_feat_dynamic_real:
        remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)

    return Chain(
        [RemoveFields(field_names=remove_field_names)]
        + (
            [SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0.0])]
            if not self.use_feat_static_cat
            else []
        )
        + [
            AsNumpyArray(field=FieldName.FEAT_STATIC_CAT, expected_ndim=1),
            AsNumpyArray(field=FieldName.TARGET, expected_ndim=1),
            # gives target the (1, T) layout
            ExpandDimArray(field=FieldName.TARGET, axis=0),
            AddObservedValuesIndicator(
                target_field=FieldName.TARGET,
                output_field=FieldName.OBSERVED_VALUES,
            ),
            # Unnormalized seasonal features
            AddTimeFeatures(
                time_features=self.issm.time_features(),
                pred_length=self.prediction_length,
                start_field=FieldName.START,
                target_field=FieldName.TARGET,
                output_field=SEASON_INDICATORS_FIELD,
            ),
            AddTimeFeatures(
                start_field=FieldName.START,
                target_field=FieldName.TARGET,
                output_field=FieldName.FEAT_TIME,
                time_features=self.time_features,
                pred_length=self.prediction_length,
            ),
            AddAgeFeature(
                target_field=FieldName.TARGET,
                output_field=FieldName.FEAT_AGE,
                pred_length=self.prediction_length,
                log_scale=True,
            ),
            VstackFeatures(
                output_field=FieldName.FEAT_TIME,
                input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
                + ([FieldName.FEAT_DYNAMIC_REAL] if self.use_feat_dynamic_real else []),
            ),
            CanonicalInstanceSplitter(
                target_field=FieldName.TARGET,
                is_pad_field=FieldName.IS_PAD,
                start_field=FieldName.START,
                forecast_start_field=FieldName.FORECAST_START,
                instance_sampler=TestSplitSampler(),
                time_series_fields=[
                    FieldName.FEAT_TIME,
                    SEASON_INDICATORS_FIELD,
                    FieldName.OBSERVED_VALUES,
                ],
                allow_target_padding=True,
                instance_length=self.past_length,
                use_prediction_features=True,
                prediction_length=self.prediction_length,
            ),
        ]
    )