Worked fine at full size...but the issue is when block_shape is not None.
probably it is running on the CPU
go back to trying self.model.predict
if 1 is indeed the case, find a way to send the model and data to the device explicitly. self.model will have no issues as it is within the strategy scope.
Worked fine at full size...but the issue is when block_shape is not None.
self.model.predictself.modelwill have no issues as it is within thestrategy scope.