(self, model_output, tile_size)
| 55 | |
| 56 | |
| 57 | def io_scale(self, model_output, tile_size): |
| 58 | # Determine the size modification happened in forward_fn |
| 59 | # We only consider the same scale on height and width. |
| 60 | io_scale = model_output.shape[2] / tile_size |
| 61 | return io_scale |
| 62 | |
| 63 | |
| 64 | def untile(self, model_output, height, width, tile_size, tile_stride, border_width, tile_device, tile_dtype): |