shallow lstm after test

main.py

@@ -121,6 +121,8 @@ if __name__ == "__main__":
         wandb.run.name = raw_run_name
         wandb.watch(network, log_freq = log_freq)
 
+    torch.autograd.set_detect_anomaly(True)
+    
     for t in range(cfg["Training"]["Epochs"]):
         print("Epoch {}".format(t))
         train(args, network, train_loader, f_loss, optimizer, device, log_interval)