| 130 | tensor_input = inputs[0]["image"] |
| 131 | |
| 132 | class WrapModel(nn.Module): |
| 133 | def __init__(self, model): |
| 134 | super().__init__() |
| 135 | if isinstance( |
| 136 | model, (nn.parallel.distributed.DistributedDataParallel, nn.DataParallel) |
| 137 | ): |
| 138 | self.model = model.module |
| 139 | else: |
| 140 | self.model = model |
| 141 | |
| 142 | def forward(self, image): |
| 143 | # jit requires the input/output to be Tensors |
| 144 | inputs = [{"image": image}] |
| 145 | outputs = self.model.forward(inputs) |
| 146 | # Only the subgraph that computes the returned tuple of tensor will be |
| 147 | # counted. So we flatten everything we found to tuple of tensors. |
| 148 | return _flatten_to_tuple(outputs) |
| 149 | |
| 150 | old_train = model.training |
| 151 | with torch.no_grad(): |
no outgoing calls
no test coverage detected