# Buggy code
def stack(
    data,
    multi_processing: bool,
    dtype: DType,
    single_process_ctx: Optional[mx.Context] = None,
):
    """Stack a list of data.
    Used when creating a single batch from list of dicts
    depending on whether multiprocessing is turned on, the batches will be
    constructed using different memory allocation techniques"""
    if isinstance(data[0], mx.nd.NDArray):
        if multi_processing:
            out = nd.empty(
                (len(data),) + data[0].shape,
                dtype=data[0].dtype,
                ctx=context.Context("cpu_shared", 0),
            )
            return mx.nd.stack(*data, out=out)
        else:
            return mx.nd.stack(*data)
    elif isinstance(data[0], np.ndarray):
        data = np.asarray(data)
        if data.dtype.kind == "f":
            data = data.astype(dtype)
        if multi_processing:
            # Workaround due to MXNet not being able to handle NDArrays with 0 in shape properly:
            if 0 in data.shape:
                return data
            return mx.nd.array(
                data, dtype=data.dtype, ctx=context.Context("cpu_shared", 0)
            )
        else:
            return mx.nd.array(data, dtype=data.dtype, ctx=single_process_ctx)
    elif isinstance(data[0], list):
        return list(
            stack(t, multi_processing, dtype, single_process_ctx) for t in zip(*data)
        )
    elif isinstance(data[0], tuple):
        return tuple(
            stack(t, multi_processing, dtype, single_process_ctx) for t in zip(*data)
        )
    elif isinstance(data[0], (pd.Timestamp, str, int, pathlib.PosixPath)):
        return data
    else:
        raise TypeError(
            f"Invalid type of data: {type(data[0])} for argument loss_function."
        )


# Traceback (TypeError)
Traceback (most recent call last):
File "/Users/stellalo/gluon-ts/issues/issue_info.py", line 25, in <module>
predictor = estimator.train(dataset)
File "/Users/stellalo/gluon-ts/src/gluonts/model/estimator.py", line 240, in train
training_data, validation_data, num_workers, num_prefetch, **kwargs
File "/Users/stellalo/gluon-ts/src/gluonts/model/estimator.py", line 219, in train_model
validation_iter=validation_data_loader,
File "/Users/stellalo/gluon-ts/src/gluonts/trainer/_base.py", line 303, in __call__
epoch_loss = loop(epoch_no, train_iter)
File "/Users/stellalo/gluon-ts/src/gluonts/trainer/_base.py", line 231, in loop
for batch_no, data_entry in enumerate(it, start=1):
File "/Users/stellalo/.virtualenvs/gluonts/lib/python3.7/site-packages/tqdm/std.py", line 1127, in __iter__
for obj in iterable:
File "/Users/stellalo/gluon-ts/src/gluonts/dataset/parallelized_loader.py", line 570, in same_process_iter
single_process_ctx=self.ctx,
File "/Users/stellalo/gluon-ts/src/gluonts/dataset/parallelized_loader.py", line 152, in default_batchify_fn
for key in data[0].keys()
File "/Users/stellalo/gluon-ts/src/gluonts/dataset/parallelized_loader.py", line 152, in <dictcomp>
for key in data[0].keys()
File "/Users/stellalo/gluon-ts/src/gluonts/dataset/parallelized_loader.py", line 133, in stack
f"Invalid type of data: {type(data[0])} for argument loss_function."
TypeError: Invalid type of data: <class 'dict'> for argument loss_function.

# Fixed code
def stack(
    data,
    multi_processing: bool,
    dtype: DType,
    single_process_ctx: Optional[mx.Context] = None,
):
    """Stack a list of data.
    Used when creating a single batch from list of dicts
    depending on whether multiprocessing is turned on, the batches will be
    constructed using different memory allocation techniques"""
    if isinstance(data[0], mx.nd.NDArray):
        if multi_processing:
            out = nd.empty(
                (len(data),) + data[0].shape,
                dtype=data[0].dtype,
                ctx=context.Context("cpu_shared", 0),
            )
            return mx.nd.stack(*data, out=out)
        else:
            return mx.nd.stack(*data)
    elif isinstance(data[0], np.ndarray):
        data = np.asarray(data)
        if data.dtype.kind == "f":
            data = data.astype(dtype)
        if multi_processing:
            # Workaround due to MXNet not being able to handle NDArrays with 0 in shape properly:
            if 0 in data.shape:
                return data
            return mx.nd.array(
                data, dtype=data.dtype, ctx=context.Context("cpu_shared", 0)
            )
        else:
            return mx.nd.array(data, dtype=data.dtype, ctx=single_process_ctx)
    elif isinstance(data[0], list):
        return list(
            stack(t, multi_processing, dtype, single_process_ctx) for t in zip(*data)
        )
    elif isinstance(data[0], tuple):
        return tuple(
            stack(t, multi_processing, dtype, single_process_ctx) for t in zip(*data)
        )

    return data