Merge branch 'master' of https://github.com/kashif/pytorch-ts

pts/model/forecast.py

@@ -215,8 +215,7 @@ class SampleForecast(Forecast):
             len(np.shape(samples)) == 2 or len(np.shape(samples)) == 3
         ), "samples should be a 2-dimensional or 3-dimensional array. Dimensions found: {}".format(
             len(np.shape(samples)))
-        self.samples = samples if (isinstance(
-            samples, np.ndarray)) else samples.numpy()
+        self.samples = samples if (isinstance(samples, np.ndarray)) else samples.cpu().numpy()
         self._sorted_samples_value = None
         self._mean = None
         self._dim = None
@@ -505,7 +504,7 @@ class DistributionForecast(Forecast):
         if self._mean is not None:
             return self._mean
         else:
-            self._mean = self.distribution.mean.numpy()
+            self._mean = self.distribution.mean.cpu().numpy()
             return self._mean
 
     @property
@@ -517,7 +516,7 @@ class DistributionForecast(Forecast):
 
     def quantile(self, level):
         level = Quantile.parse(level).value
-        q = self.distribution.icdf(torch.tensor([level])).numpy()
+        q = self.distribution.icdf(torch.tensor([level])).cpu().numpy()
         return q
 
     def to_sample_forecast(self, num_samples: int = 200) -> SampleForecast:

pts/model/forecast_generator.py

@@ -95,7 +95,7 @@ class QuantileForecastGenerator(ForecastGenerator):
                  num_samples: Optional[int], **kwargs) -> Iterator[Forecast]:
         for batch in inference_data_loader:
             inputs = [batch[k] for k in input_names]
-            outputs = prediction_net(*inputs).numpy()
+            outputs = prediction_net(*inputs).cpu().numpy()
             if output_transform is not None:
                 outputs = output_transform(batch, outputs)
 
@@ -120,14 +120,14 @@ class SampleForecastGenerator(ForecastGenerator):
                  num_samples: Optional[int], **kwargs) -> Iterator[Forecast]:
         for batch in inference_data_loader:
             inputs = [batch[k] for k in input_names]
-            outputs = prediction_net(*inputs).numpy()
+            outputs = prediction_net(*inputs).cpu().numpy()
             if output_transform is not None:
                 outputs = output_transform(batch, outputs)
             if num_samples:
                 num_collected_samples = outputs[0].shape[0]
                 collected_samples = [outputs]
                 while num_collected_samples < num_samples:
-                    outputs = prediction_net(*inputs).numpy()
+                    outputs = prediction_net(*inputs).cpu().numpy()
                     if output_transform is not None:
                         outputs = output_transform(batch, outputs)
                     collected_samples.append(outputs)

pts/model/predictor.py

@@ -59,11 +59,12 @@ class PTSPredictor(Predictor):
             device=self.device,
             dtype=self.dtype,
         )
-        yield from self.forecast_generator(
-            inference_data_loader=inference_data_loader,
-            prediction_net=self.prediction_net,
-            input_names=self.input_names,
-            freq=self.freq,
-            output_transform=self.output_transform,
-            num_samples=num_samples,
-        )
+        with torch.no_grad():
+            yield from self.forecast_generator(
+                inference_data_loader=inference_data_loader,
+                prediction_net=self.prediction_net,
+                input_names=self.input_names,
+                freq=self.freq,
+                output_transform=self.output_transform,
+                num_samples=num_samples,
+            )