Make explicit cast to np.float64 in nns

tinkoff-ai · Jan 19, 2023 · 38385ef · 38385ef
1 parent bfbd4b1
commit 38385ef
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/etna/models/base.py b/etna/models/base.py
@@ -632,8 +632,7 @@ def forecast(self, ts: "TSDataset", prediction_size: int) -> "TSDataset":
         future_ts = ts.tsdataset_idx_slice(start_idx=self.encoder_length, end_idx=self.encoder_length + prediction_size)
         for (segment, feature_nm), value in predictions.items():
             # we don't want to change dtype after assignment, but there can happen cast to float32
-            dtype = future_ts.df.loc[:, pd.IndexSlice[segment, feature_nm]]
-            future_ts.df.loc[:, pd.IndexSlice[segment, feature_nm]] = value[:prediction_size, :].astype(dtype)
+            future_ts.df.loc[:, pd.IndexSlice[segment, feature_nm]] = value[:prediction_size, :].astype(np.float64)
 
         future_ts.inverse_transform()