| # Copyright (c) Facebook, Inc. and its affiliates. | |
| # | |
| # This source code is licensed under the MIT license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| import torch | |
| from .task import Task | |
| class MILNCETask(Task): | |
| def reshape_subsample(self, sample): | |
| if ( | |
| hasattr(self.config.dataset, "subsampling") | |
| and self.config.dataset.subsampling is not None | |
| and self.config.dataset.subsampling > 1 | |
| ): | |
| for key in sample: | |
| if torch.is_tensor(sample[key]): | |
| tensor = self.flat_subsample(sample[key]) | |
| if key in ["caps", "cmasks"]: | |
| size = tensor.size() | |
| batch_size = size[0] * size[1] | |
| expanded_size = (batch_size,) + size[2:] | |
| tensor = tensor.view(expanded_size) | |
| sample[key] = tensor | |
| return sample | |