File tree Expand file tree Collapse file tree 2 files changed +0
-8
lines changed Expand file tree Collapse file tree 2 files changed +0
-8
lines changed Original file line number Diff line number Diff line change @@ -387,10 +387,6 @@ def train_init(
387
387
aug_rng = jax .random .PRNGKey (self .config ["seed" ])
388
388
batch = aug_fn (aug_rng , batch )
389
389
390
- # check image size
391
- image_shape = self .dataset_info .image_spatial_shape
392
- chex .assert_equal (batch [IMAGE ].shape [1 :- 1 ], image_shape )
393
-
394
390
# init train state on cpu first
395
391
dtype = get_half_precision_dtype (self .config .half_precision )
396
392
model = instantiate (self .config .task .model , dtype = dtype )
Original file line number Diff line number Diff line change @@ -261,10 +261,6 @@ def train_init(
261
261
aug_rng = jax .random .PRNGKey (self .config ["seed" ])
262
262
batch = aug_fn (aug_rng , batch )
263
263
264
- # check image size
265
- image_shape = self .dataset_info .image_spatial_shape
266
- chex .assert_equal (batch [IMAGE ].shape [1 :- 1 ], image_shape )
267
-
268
264
# init train state on cpu first
269
265
dtype = get_half_precision_dtype (self .config .half_precision )
270
266
model = instantiate (self .config .task .model , dtype = dtype )
You can’t perform that action at this time.
0 commit comments