From 4241bff32aee728f4d2b57f52c667e15ece3e33c Mon Sep 17 00:00:00 2001
From: Yinon Polak <inon51@gmail.com>
Date: Mon, 6 Mar 2023 20:15:36 +0200
Subject: [PATCH] type hints fixes

---
 freqtrade/freqai/base_models/PyTorchModelTrainer.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/freqtrade/freqai/base_models/PyTorchModelTrainer.py b/freqtrade/freqai/base_models/PyTorchModelTrainer.py
index 02ff35085..fc0a7600e 100644
--- a/freqtrade/freqai/base_models/PyTorchModelTrainer.py
+++ b/freqtrade/freqai/base_models/PyTorchModelTrainer.py
@@ -1,6 +1,7 @@
 import logging
 from pathlib import Path
 from typing import Dict
+from torch.optim import Optimizer
 
 import torch
 import torch.nn as nn
@@ -15,7 +16,7 @@ class PyTorchModelTrainer:
     def __init__(
             self,
             model: nn.Module,
-            optimizer: nn.Module,
+            optimizer: Optimizer,
             criterion: nn.Module,
             device: str,
             batch_size: int,