dataset type fix

This commit is contained in:
shahules786 2022-09-24 12:46:40 +05:30
parent 48d5f9c21e
commit 71b98ba67c
1 changed files with 2 additions and 2 deletions

View File

@ -13,7 +13,7 @@ from pathlib import Path
from enhancer import __version__ from enhancer import __version__
from enhancer.data.dataset import Dataset from enhancer.data.dataset import EnhancerDataset
from enhancer.utils.io import Audio from enhancer.utils.io import Audio
from enhancer.utils.loss import Avergeloss from enhancer.utils.loss import Avergeloss
from enhancer.inference import Inference from enhancer.inference import Inference
@ -29,7 +29,7 @@ class Model(pl.LightningModule):
num_channels:int=1, num_channels:int=1,
sampling_rate:int=16000, sampling_rate:int=16000,
lr:float=1e-3, lr:float=1e-3,
dataset:Optional[Dataset]=None, dataset:Optional[EnhancerDataset]=None,
duration:Optional[float]=None, duration:Optional[float]=None,
loss: Union[str, List] = "mse", loss: Union[str, List] = "mse",
metric:Union[str,List] = "mse" metric:Union[str,List] = "mse"