fix some tests

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,

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)
 

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

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(