update code

main.py

@@ -4,7 +4,7 @@ import os
 import pytorch_lightning as pl
 import soundfile as sf
 import torch
-from pytorch_lightning.callbacks import ModelCheckpoint, StochasticWeightAveraging
+from pytorch_lightning.callbacks import ModelCheckpoint
 from pytorch_lightning.utilities.model_summary import summarize
 from torch.utils.data import DataLoader
 
@@ -66,7 +66,7 @@ def train():
                          gpus=len(gpus),
                          max_epochs=CONFIG.TRAIN.epochs,
                          accelerator="gpu" if len(gpus) > 1 else None,
-                         callbacks=[checkpoint_callback, StochasticWeightAveraging(swa_lrs=1e-2)]
+                         callbacks=[checkpoint_callback]
                          )
 
     print(model.hparams)

models/blocks.py

@@ -117,7 +117,7 @@ class Predictor(pl.LightningModule):  # mel
         fb = librosa.filters.mel(sr=sr, n_fft=self.window_size, n_mels=self.n_mels)[:, 1:]
         self.fb = torch.from_numpy(fb).unsqueeze(0).unsqueeze(0)
         self.lstm = nn.LSTM(input_size=self.n_mels, hidden_size=self.lstm_dim, bidirectional=False,
-                            num_layers=self.lstm_layers)
+                            num_layers=self.lstm_layers, batch_first=True)
         self.expand_dim = nn.Linear(self.lstm_dim, self.n_mels)
         self.inv_mel = nn.Linear(self.n_mels, self.hop_size)
 

models/frn.py

@@ -66,7 +66,7 @@ class PLCModel(pl.LightningModule):
 
         x = x.permute(3, 0, 1, 2).unsqueeze(-1)
         prev_mag = torch.zeros((B, 1, F, 1), device=x.device)
-        predictor_state = torch.zeros((2, self.predictor.lstm_layers, 1, self.predictor.lstm_dim), device=x.device)
+        predictor_state = torch.zeros((2, self.predictor.lstm_layers, B, self.predictor.lstm_dim), device=x.device)
         mlp_state = torch.zeros((self.encoder.depth, 2, 1, B, self.encoder.dim), device=x.device)
         result = []
         for step in x:
@@ -201,7 +201,7 @@ class OnnxWrapper(pl.LightningModule):
         super().__init__(*args, **kwargs)
         self.model = model
         batch_size = 1
-        pred_states = torch.zeros((2, 1, 1, model.predictor.lstm_dim))
+        pred_states = torch.zeros((2, 1, batch_size, model.predictor.lstm_dim))
         mlp_states = torch.zeros((model.encoder.depth, 2, 1, batch_size, model.encoder.dim))
         mag = torch.zeros((batch_size, 1, model.hop_size, 1))
         x = torch.randn(batch_size, model.hop_size + 1, 2)