CI failing randomly with NaN arrays in state_dict_io test

[jdeschamps]

Describe the bug

A problem that had disappeared is not present again in #365: some tests randomly fail with NaN arrays.

 ______________________________ test_state_dict_io ______________________________

tmp_path = PosixPath('/tmp/pytest-of-runner/pytest-0/test_state_dict_io0')
ordered_array = <function ordered_array.<locals>._ordered_array at 0x7f5f18672020>
pre_trained = PosixPath('/tmp/pytest-of-runner/pytest-0/test_state_dict_io0/checkpoints/last.ckpt')

    def test_state_dict_io(tmp_path, ordered_array, pre_trained):
        """Test exporting and loading a state dict."""
        # training data
        train_array = ordered_array((32, 32))
        path = tmp_path / "model.pth"
    
        # instantiate CAREamist
        careamist = CAREamist(source=pre_trained, work_dir=tmp_path)
    
        # predict (no tiling and no tta)
        predicted_output = careamist.predict(train_array, tta_transforms=False)
        predicted = np.concatenate(predicted_output, axis=0)
    
        # save model
        _export_state_dict(careamist.model, path)
        assert path.exists()
    
        # load model
        _load_state_dict(careamist.model, path)
    
        # predict (no tiling and no tta)
        predicted_loaded = careamist.predict(train_array, tta_transforms=False)

>       assert (predicted_loaded == predicted).all()
E       assert False
E        +  where False = <built-in method all of numpy.ndarray object at 0x7f5f186ca6d0>()
E        +    where <built-in method all of numpy.ndarray object at 0x7f5f186ca6d0> = [array([[[[na...type=float32)] == array([[[[nan...dtype=float32)
E             
E             Full diff:
E             + [
E             - array([[[[nan, nan, nan, ..., nan, nan, nan],
E             +     array([[[[nan, nan, nan, ..., nan, nan, nan],
E             ? ++++
E                        [nan, nan, nan, ..., nan, nan, nan],
E                        [nan, nan, nan, ..., nan, nan, nan],
E                        ...,
E                        [nan, nan, nan, ..., nan, nan, nan],
E                        [nan, nan, nan, ..., nan, nan, nan],
E             -          [nan, nan, nan, ..., nan, nan, nan]]]], dtype=float32)
E             +          [nan, nan, nan, ..., nan, nan, nan]]]], dtype=float32),
E             ?                                                                +
E             + ].all

tests/model_io/test_bmz_io.py:32: AssertionError

If we re-run the tests, then they often pass again. It is very inconsistent.

Anybody remembers if we ever figured it out? Did this happen locally as well? Or is this PR changing something that triggers it more?

@CatEek @melisande-c @veegalinova

[melisande-c]

Did we find (only empirically) that np.testing.assert_array_equal was more stable?

[melisande-c]

But also I can't see any of the tests failing for that reason?

[jdeschamps]

But also I can't see any of the tests failing for that reason?

Yeah because I manually run them again so that they pass.... I will edit the main post with this info.