Changed hardcoded type "float32" to dtype (models.py)

PiperOrigin-RevId: 321632088

Author: Flax Team

examples/imagenet/models.py

@@ -78,7 +78,7 @@ def apply(self, x, num_classes, num_filters=64, num_layers=50,
                           dtype=dtype)
     x = jnp.mean(x, axis=(1, 2))
     x = nn.Dense(x, num_classes, dtype=dtype)
-    x = jnp.asarray(x, jnp.float32)
+    x = jnp.asarray(x, dtype)
     x = nn.log_softmax(x)
     return x