Add batch_stats verification in test_training_mode

Co-authored-by: or4k2l <219930442+or4k2l@users.noreply.github.com>

Author: Copilot

tests/test_model.py

@@ -54,6 +54,9 @@ def test_training_mode(self, model, test_input):
         # Extract logits and updated batch_stats
         logits, updated_vars = output
 
+        # Verify batch_stats is present (model uses BatchNorm)
+        assert 'batch_stats' in updated_vars
+        
         # Should produce valid output
         assert logits.shape == (4, 10)
         assert not jnp.any(jnp.isnan(logits))