(self, sample, tile_size=64, tile_stride=32)
| 39 | self.conv_out = torch.nn.Conv2d(512, 8, kernel_size=3, padding=1) |
| 40 | |
| 41 | def tiled_forward(self, sample, tile_size=64, tile_stride=32): |
| 42 | hidden_states = TileWorker().tiled_forward( |
| 43 | lambda x: self.forward(x), |
| 44 | sample, |
| 45 | tile_size, |
| 46 | tile_stride, |
| 47 | tile_device=sample.device, |
| 48 | tile_dtype=sample.dtype |
| 49 | ) |
| 50 | return hidden_states |
| 51 | |
| 52 | def forward(self, sample, tiled=False, tile_size=64, tile_stride=32, **kwargs): |
| 53 | original_dtype = sample.dtype |
no test coverage detected