Flaky tests when using `model_probably_good`

Tests using `model_probably_good` are flaky and sometimes fail.

```
..................................................F....
======================================================================
FAIL: test_SkipLinear_different_dim (test_nn.TestLinear.test_SkipLinear_different_dim)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/Users/distiller/project/tests/test_nn.py", line 86, in test_SkipLinear_different_dim
self.assertTrue(model_probably_good(model, (din,), (dout, )))
AssertionError: False is not true
```

Should look into either switching them out for non-random tests or make sure that they pass with much more certainty.