Merge branch 'main' into more_forward_intermediates

Author: brianhou0208

tests/test_models.py

@@ -511,8 +511,9 @@ def test_model_forward_intermediates(model_name, batch_size):
     spatial_axis = get_spatial_dim(output_fmt)
     import math
 
+    inpt = torch.randn((batch_size, *input_size))
     output, intermediates = model.forward_intermediates(
-        torch.randn((batch_size, *input_size)),
+        inpt,
         output_fmt=output_fmt,
     )
     assert len(expected_channels) == len(intermediates)
@@ -524,6 +525,9 @@ def test_model_forward_intermediates(model_name, batch_size):
         assert o.shape[0] == batch_size
         assert not torch.isnan(o).any()
 
+    output2 = model.forward_features(inpt)
+    assert torch.allclose(output, output2)
+
 
 def _create_fx_model(model, train=False):
     # This block of code does a bit of juggling to handle any case where there are multiple outputs in train mode