fixing linting errors

mk314k · Nov 21, 2023 · bf26541 · bf26541
1 parent b9dc569
commit bf26541
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -47,7 +47,7 @@ def vae_loss(recon_x, label):
 def train( # pylint: disable=too-many-locals
  train_x:torch.Tensor,
  train_y:torch.Tensor,
- models,
+ models:tuple[nn.Module],
  optims,
  num_epochs=10
 ):