# Buggy code
def get_issm_coeff(
    self,
    seasonal_indicators: Tensor,  # (batch_size, time_length)
) -> Tuple[Tensor, Tensor, Tensor]:
    F = getF(seasonal_indicators)
    emission_coeff_ls, transition_coeff_ls, innovation_coeff_ls = zip(
        self.nonseasonal_issm.get_issm_coeff(seasonal_indicators),
        *[
            issm.get_issm_coeff(
                seasonal_indicators.slice_axis(axis=-1, begin=ix, end=ix + 1)
            )
            for ix, issm in enumerate(self.seasonal_issms)
        ],
    )

    # stack emission and innovation coefficients
    emission_coeff = F.concat(*emission_coeff_ls, dim=-1)

    innovation_coeff = F.concat(*innovation_coeff_ls, dim=-1)

    # transition coefficient is block diagonal!
    transition_coeff = _make_block_diagonal(transition_coeff_ls)

    return emission_coeff, transition_coeff, innovation_coeff


# 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 get_issm_coeff(
    self,
    features: Tensor,  # (batch_size, time_length, num_features)
) -> Tuple[Tensor, Tensor, Tensor]:
    F = getF(features)
    emission_coeff_ls, transition_coeff_ls, innovation_coeff_ls = zip(
        *[
            issm.get_issm_coeff(features.slice_axis(axis=-1, begin=ix, end=ix + 1))
            for ix, issm in enumerate([self.nonseasonal_issm] + self.seasonal_issms)
        ],
    )

    # stack emission and innovation coefficients
    emission_coeff = F.concat(*emission_coeff_ls, dim=-1)

    innovation_coeff = F.concat(*innovation_coeff_ls, dim=-1)

    # transition coefficient is block diagonal!
    transition_coeff = _make_block_diagonal(transition_coeff_ls)

    return emission_coeff, transition_coeff, innovation_coeff