modified: test_models.py

nkern · Feb 7, 2024 · 5c2d90e · 5c2d90e
1 parent 442fb68
commit 5c2d90e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/py21cmnet/tests/test_models.py b/py21cmnet/tests/test_models.py
@@ -98,7 +98,7 @@ def test_autoencoder():
             info = utils.train(model, dl, torch.nn.MSELoss(reduction='mean'), torch.optim.Adam,
                                optim_kwargs=dict(lr=0.1), Nepochs=3, verbose=True)
             # assert loss decreases
-            assert (np.diff(info['train_loss'].detach().numpy()) < 0).all()
+            assert (np.diff(torch.stack(info['train_loss']).detach().numpy()) < 0).all()
 
             # pred = model(X)
             # import matplotlib.pyplot as plt;