type hints fixes

2023-03-06 19:14:54 +02:00
parent 125085fbaf
commit 8acdd0b47c
2 changed files with 3 additions and 2 deletions
--- a/freqtrade/freqai/prediction_models/PyTorchMLPModel.py
+++ b/freqtrade/freqai/prediction_models/PyTorchMLPModel.py
@@ -3,6 +3,7 @@ import logging

 import torch
 import torch.nn as nn
+from torch import Tensor

 logger = logging.getLogger(__name__)

@@ -16,7 +17,7 @@ class PyTorchMLPModel(nn.Module):
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(p=0.2)

-    def forward(self, x: torch.tensor) -> torch.tensor:
+    def forward(self, x: Tensor) -> Tensor:
        x = self.relu(self.input_layer(x))
        x = self.dropout(x)
        x = self.relu(self.hidden_layer(x))