# Buggy code
def __init__(
    self,
    dataset: Dataset,
    transformation: Transformation,
    cyclic: bool,
    is_train: bool,
    batch_size: int,
    shuffle: bool = False,
    batchify_fn: Callable = None,
    ctx: mx.Context = None,
    dtype: DType = np.float32,
    num_prefetch: Optional[int] = None,
    num_workers: Optional[int] = None,
):
    # Some windows error with the ForkingPickler prevents usage currently:
    if sys.platform == "win32":
        logging.warning(
            "You have set `num_workers` for to a non zero value, "
            "however, currently multiprocessing is not supported on windows."
        )
        num_workers = 0

    self.dataset = dataset
    self.dataset_len = None
    if isinstance(dataset, Sized):
        assert isinstance(dataset, Sized)
        self.dataset_len = len(dataset)
    else:
        self.dataset_len = len(list(dataset))
    # indicates that we want to cycle through the dataset
    self.cyclic = cyclic
    # indicates the current cycle, needed for resetting iterators at each cycle
    self.cycle_num = 0

    self.dtype = dtype
    self.is_train = is_train
    self.transformation = transformation
    self.ctx = ctx
    self.batch_size = batch_size
    self.shuffle = shuffle

    assert num_workers is None or num_workers <= self.dataset_len, (
        "Cannot have more workers than dataset entries currently."
    )

    # TODO: switch to default multiprocessing.cpu_count() here
    default_num_workers = 0
    self.num_workers = max(
        0,
        num_workers
        if num_workers is not None
        else min(self.dataset_len, default_num_workers),
    )
    self.num_prefetch = max(
        0,
        num_prefetch if num_prefetch is not None else 2 * self.num_workers,
    )
    self.worker_pool = None
    # In order to set unique IDs to workers:
    self.worker_manager = None
    self.worker_id_queue = None
    # In order to recycle unused but pre-calculated batches from last epoch for training:
    self.multi_worker_cache = None

    if self.num_workers > 0:
        # generate unique ids for processes
        self.worker_manager = multiprocessing.Manager()
        self.worker_id_queue = self.worker_manager.Queue()
        for i in range(self.num_workers):
            self.worker_id_queue.put(i)

        self.worker_pool = multiprocessing.get_context("spawn").Pool(
            self.num_workers,
            initializer=_worker_initializer,
            initargs=[
                self.dataset,
                self.dataset_len,
                self.num_workers,
                self.transformation,
                self.cyclic,
                self.worker_id_queue,
            ],
        )

    if batchify_fn is None:
        self.batchify_fn = default_batchify_fn
    else:
        self.batchify_fn = batchify_fn


# Traceback (mxnet.base.MXNetError)
Traceback (most recent call last):
File "/usr/local/lib/python3.6/runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "/usr/local/lib/python3.6/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/opt/ml/code/train.py", line 137, in <module>
forecasts, tss, agg_metrics, item_metrics = evaluate(predictor, dataset.test)
File "/opt/ml/code/utils.py", line 30, in evaluate
forecasts = list(forecast_it)
File "/usr/local/lib/python3.6/site-packages/gluonts/model/predictor.py", line 317, in predict
num_samples=num_samples,
File "/usr/local/lib/python3.6/site-packages/gluonts/model/forecast_generator.py", line 195, in __call__
for batch in inference_data_loader:
File "/usr/local/lib/python3.6/site-packages/gluonts/dataset/loader.py", line 101, in __iter__
yield from self.parallel_data_loader
File "/usr/local/lib/python3.6/site-packages/gluonts/dataset/parallelized_loader.py", line 572, in same_process_iter
for k, v in batch.items()
File "/usr/local/lib/python3.6/site-packages/gluonts/dataset/parallelized_loader.py", line 572, in <dictcomp>
for k, v in batch.items()
File "/usr/local/lib/python3.6/site-packages/mxnet/ndarray/ndarray.py", line 2765, in as_in_context
return self.copyto(context)
File "/usr/local/lib/python3.6/site-packages/mxnet/ndarray/ndarray.py", line 2632, in copyto
return _internal._copyto(self, out=hret)
File "<string>", line 27, in _copyto
File "/usr/local/lib/python3.6/site-packages/mxnet/_ctypes/ndarray.py", line 107, in _imperative_invoke
ctypes.byref(out_stypes)))
File "/usr/local/lib/python3.6/site-packages/mxnet/base.py", line 255, in check_call
raise MXNetError(py_str(_LIB.MXGetLastError()))
mxnet.base.MXNetError: [18:26:25] src/imperative/./imperative_utils.h:146: Operator _copyto inferring shapes failed.
input shapes:
[1,321,-1]
output shapes:
[1,321,-1]
operator attributes:

Stack trace:
[bt] (0) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x6dfb0b) [0x7f99a9e28b0b]
[bt] (1) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(mxnet::imperative::SetShapeType(mxnet::Context const&amp;, nnvm::NodeAttrs const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, mxnet::DispatchMode*)+0x363b) [0x7f99ad0da57b]
[bt] (2) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(mxnet::Imperative::Invoke(mxnet::Context const&amp;, nnvm::NodeAttrs const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;, std::vector<mxnet::NDArray*, std::allocator<mxnet::NDArray*> > const&amp;)+0x1db) [0x7f99ad0e24bb]
[bt] (3) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x3847f3f) [0x7f99acf90f3f]
[bt] (4) /usr/local/lib/python3.6/site-packages/mxnet/libmxnet.so(MXImperativeInvokeEx+0x62) [0x7f99acf91502]
[bt] (5) /usr/local/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(ffi_call_unix64+0x4c) [0x7f9a01ccf4e6]
[bt] (6) /usr/local/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(ffi_call+0x3f1) [0x7f9a01cce241]
[bt] (7) /usr/local/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(_ctypes_callproc+0x2cf) [0x7f9a01cc58ff]
[2020-04-02 18:26:25.315 ip-10-0-168-31.us-west-2.compute.internal:213 INFO utils.py:25] The end of training job file will not be written for jobs running under SageMaker.
[bt] (8) /usr/local/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(+0x9829) [0x7f9a01cbc829]

# Fixed code
def __init__(
    self,
    dataset: Dataset,
    transformation: Transformation,
    cyclic: bool,
    is_train: bool,
    batch_size: int,
    shuffle: bool = False,
    batchify_fn: Callable = None,
    ctx: mx.Context = None,
    dtype: DType = np.float32,
    num_prefetch: Optional[int] = None,
    num_workers: Optional[int] = None,
):
    # Some windows error with the ForkingPickler prevents usage currently:
    if sys.platform == "win32":
        logging.warning(
            "You have set `num_workers` for to a non zero value, "
            "however, currently multiprocessing is not supported on windows."
        )
        num_workers = 0

    self.dataset = dataset
    self.dataset_len = None
    if isinstance(dataset, Sized):
        assert isinstance(dataset, Sized)
        self.dataset_len = len(dataset)
    else:
        self.dataset_len = len(list(dataset))
    # indicates that we want to cycle through the dataset
    self.cyclic = cyclic
    # indicates the current cycle, needed for resetting iterators at each cycle
    self.cycle_num = 0

    self.dtype = dtype
    self.is_train = is_train
    self.transformation = transformation
    self.ctx = ctx
    self.batch_size = batch_size
    self.shuffle = shuffle

    assert num_workers is None or num_workers <= self.dataset_len, (
        "Cannot have more workers than dataset entries currently."
    )

    # TODO: switch to default multiprocessing.cpu_count() here
    default_num_workers = 0
    self.num_workers = max(
        0,
        num_workers
        if num_workers is not None
        else min(self.dataset_len, default_num_workers),
    )
    self.num_prefetch = max(
        0,
        num_prefetch if num_prefetch is not None else 2 * self.num_workers,
    )
    self.worker_pool = None
    # In order to set unique IDs to workers:
    self.worker_manager = None
    self.worker_id_queue = None
    # In order to recycle unused but pre-calculated batches from last epoch for training:
    self.multi_worker_cache = None

    if self.num_workers > 0:
        # generate unique ids for processes
        self.worker_manager = multiprocessing.Manager()
        self.worker_id_queue = self.worker_manager.Queue()
        for i in range(self.num_workers):
            self.worker_id_queue.put(i)

        self.worker_pool = multiprocessing.get_context("spawn").Pool(
            self.num_workers,
            initializer=_worker_initializer,
            initargs=[
                self.dataset,
                self.transformation,
                self.num_workers,
                self.worker_id_queue,
            ],
        )

    if batchify_fn is None:
        self.batchify_fn = default_batchify_fn
    else:
        self.batchify_fn = batchify_fn