From d5577a2c9e43573a93d8166976e165537de1bb51 Mon Sep 17 00:00:00 2001
From: "Dr. Kashif Rasul" <kashif.rasul@gmail.com>
Date: Wed, 30 Dec 2020 19:17:54 +0100
Subject: [PATCH] fix some tests

---
 pts/model/estimator.py                      |  4 ++--
 pts/modules/scaler.py                       |  4 ++++
 test/model/deepar/test_auxillary_outputs.py |  9 ++++++---
 test/modules/test_distribution_output.py    | 14 +++++++++-----
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/pts/model/estimator.py b/pts/model/estimator.py
index 8962993..f54967b 100644
--- a/pts/model/estimator.py
+++ b/pts/model/estimator.py
@@ -79,7 +79,7 @@ class PyTorchEstimator(Estimator):
         self,
         training_data: Dataset,
         validation_data: Optional[Dataset] = None,
-        num_workers: Optional[int] = None,
+        num_workers: int = 0,
         prefetch_factor: int = 2,
         shuffle_buffer_length: Optional[int] = None,
         **kwargs,
@@ -138,7 +138,7 @@ class PyTorchEstimator(Estimator):
         self,
         training_data: Dataset,
         validation_data: Optional[Dataset] = None,
-        num_workers: Optional[int] = None,
+        num_workers: int = 0,
         prefetch_factor: int = 2,
         shuffle_buffer_length: Optional[int] = None,
         **kwargs,
diff --git a/pts/modules/scaler.py b/pts/modules/scaler.py
index 7a80e3a..3755592 100644
--- a/pts/modules/scaler.py
+++ b/pts/modules/scaler.py
@@ -4,6 +4,8 @@ from typing import Tuple
 import torch
 import torch.nn as nn
 
+from gluonts.core.component import validated
+
 
 class Scaler(ABC, nn.Module):
     def __init__(self, keepdim: bool = False, time_first: bool = True):
@@ -68,6 +70,7 @@ class MeanScaler(Scaler):
         default scale that is used if the time series has only zeros.
     """
 
+    @validated()
     def __init__(self, minimum_scale: float = 1e-10, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.register_buffer("minimum_scale", torch.tensor(minimum_scale))
@@ -111,6 +114,7 @@ class NOPScaler(Scaler):
     no scaling is applied upon calling the ``NOPScaler``.
     """
 
+    @validated()
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
diff --git a/test/model/deepar/test_auxillary_outputs.py b/test/model/deepar/test_auxillary_outputs.py
index 14c36c4..413fb56 100644
--- a/test/model/deepar/test_auxillary_outputs.py
+++ b/test/model/deepar/test_auxillary_outputs.py
@@ -15,8 +15,11 @@ from itertools import islice
 
 import torch
 
+from gluonts.dataset.artificial import constant_dataset
+from gluonts.dataset.loader import TrainDataLoader
+from gluonts.torch.batchify import batchify 
+
 from pts import Trainer
-from pts.dataset import constant_dataset, TrainDataLoader
 from pts.model import get_module_forward_input_names
 from pts.model.deepar import DeepAREstimator
 from pts.modules import StudentTOutput
@@ -46,11 +49,11 @@ def test_distribution():
     num_samples = 3
 
     training_data_loader = TrainDataLoader(
-        dataset=train_ds,
+        train_ds,
         transform=train_output.transformation,
         batch_size=batch_size,
         num_batches_per_epoch=estimator.trainer.num_batches_per_epoch,
-        device=torch.device("cpu"),
+        stack_fn=batchify,
     )
 
     seq_len = 2 * ds_info.prediction_length
diff --git a/test/modules/test_distribution_output.py b/test/modules/test_distribution_output.py
index 6a818b3..00605d9 100644
--- a/test/modules/test_distribution_output.py
+++ b/test/modules/test_distribution_output.py
@@ -127,9 +127,7 @@ def test_neg_binomial(total_count_logit: Tuple[float, float]) -> None:
     total_counts = torch.zeros((NUM_SAMPLES,)) + total_count
     logits = torch.zeros((NUM_SAMPLES,)) + logit
 
-    neg_bin_distr = NegativeBinomial(
-        total_count=total_counts, logits=logits
-    )
+    neg_bin_distr = NegativeBinomial(total_count=total_counts, logits=logits)
     samples = neg_bin_distr.sample()
 
     init_biases = [
@@ -138,7 +136,10 @@ def test_neg_binomial(total_count_logit: Tuple[float, float]) -> None:
     ]
 
     total_count_hat, logit_hat = maximum_likelihood_estimate_sgd(
-        NegativeBinomialOutput(), samples, init_biases=init_biases, num_epochs=15,
+        NegativeBinomialOutput(),
+        samples,
+        init_biases=init_biases,
+        num_epochs=15,
     )
 
     assert (
@@ -201,7 +202,10 @@ def test_independent_normal() -> None:
     samples = distr.sample((num_samples,))
 
     loc_hat, diag_hat = maximum_likelihood_estimate_sgd(
-        NormalOutput(dim=dim), samples, learning_rate=0.01, num_epochs=10,
+        NormalOutput(dim=dim),
+        samples,
+        learning_rate=0.01,
+        num_epochs=10,
     )
 
     distr = Independent(