Error reproducing results with efficientnet on SVHN dataset, the size of input is more then 224, but image is 32x32

[MotiBaadror]

I0722 07:12:25.574716 140024631277440 tpu_client.py:54] Starting the local TPU driver.
I0722 07:12:25.575148 140024631277440 xla_bridge.py:214] Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker: local://
I0722 07:12:25.784465 140024631277440 xla_bridge.py:214] Unable to initialize backend 'tpu': Invalid argument: TpuPlatform is not available.
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:365: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
"jax.host_count has been renamed to jax.process_count. This alias "
I0722 07:12:26.047336 140024631277440 train.py:106] Total batch size: 16 (16 x 1 replicas)
I0722 07:12:26.048122 140024631277440 dataset_info.py:362] Load dataset info from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.050460 140024631277440 dataset_builder.py:323] Reusing dataset svhn_cropped (/root/tensorflow_datasets/svhn_cropped/3.0.0)
I0722 07:12:26.050613 140024631277440 dataset_builder.py:529] Constructing tf.data.Dataset for split train, from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.170840 140024631277440 dataset_info.py:362] Load dataset info from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.172739 140024631277440 dataset_builder.py:323] Reusing dataset svhn_cropped (/root/tensorflow_datasets/svhn_cropped/3.0.0)
I0722 07:12:26.172901 140024631277440 dataset_builder.py:529] Constructing tf.data.Dataset for split extra, from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.244570 140024631277440 dataset_info.py:362] Load dataset info from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.246468 140024631277440 dataset_builder.py:323] Reusing dataset svhn_cropped (/root/tensorflow_datasets/svhn_cropped/3.0.0)
I0722 07:12:26.246642 140024631277440 dataset_builder.py:529] Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/svhn_cropped/3.0.0
I0722 07:12:26.322749 140024631277440 dataset_source.py:315] Used test set instead of validation set.
Traceback (most recent call last):
File "/usr/lib/python3.7/runpy.py", line 193, in _run_module_as_main
"main", mod_spec)
File "/usr/lib/python3.7/runpy.py", line 85, in _run_code
exec(code, run_globals)
File "/content/sam/sam_jax/train.py", line 177, in
app.run(main)
File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 303, in run
_run_main(main, args)
File "/usr/local/lib/python3.7/dist-packages/absl/app.py", line 251, in _run_main
sys.exit(main(argv))
File "/content/sam/sam_jax/train.py", line 162, in main
num_classes)
File "/content/sam/sam_jax/imagenet_models/load_model.py", line 134, in get_model
module)
File "/content/sam/sam_jax/imagenet_models/load_model.py", line 62, in create_image_model
prng_key, [(input_shape, jnp.float32)])
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 220, in wrapper
return super_fn(*args, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 220, in wrapper
return super_fn(*args, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 512, in init_by_shape
return jax_utils.partial_eval_by_shape(lazy_init, input_specs)
File "/usr/local/lib/python3.7/dist-packages/flax/jax_utils.py", line 119, in partial_eval_by_shape
_, out_pvals, _ = pe.trace_to_jaxpr(f_flat, in_pvals)
File "/usr/local/lib/python3.7/dist-packages/jax/interpreters/partial_eval.py", line 505, in trace_to_jaxpr
jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
File "/usr/local/lib/python3.7/dist-packages/jax/linear_util.py", line 166, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/lib/python3.7/dist-packages/flax/jax_utils.py", line 111, in
f = lambda *inputs: fn(*inputs, *args, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 509, in lazy_init
return init_fn()
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 503, in init_fn
return cls.init(_rng, *(inputs + args), name=name, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 220, in wrapper
return super_fn(*args, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 220, in wrapper
return super_fn(*args, **kwargs)
File "/usr/local/lib/python3.7/dist-packages/flax/nn/base.py", line 465, in init
y = instance.apply(*args, **kwargs)
File "/content/sam/sam_jax/efficientnet/efficientnet.py", line 597, in apply
'but got input of resolution {}'.format(x.shape[1:3]))
ValueError: Wrong input size. Model was expecting resolution (224, 224) but got input of resolution (32, 32)