(self, input_batch: NestedTensor)
| 167 | return y.mean(), {} |
| 168 | |
| 169 | def predict(self, input_batch: NestedTensor) -> Tensor: |
| 170 | x = input_batch["x"] |
| 171 | with child_context("input_stats", module=self): |
| 172 | self.add_module_output("x_mean", x.mean()) |
| 173 | self.add_module_output("x_max", x.max()) |
| 174 | self.predict_dtypes.append(x.dtype) |
| 175 | return self.linear(x) |
| 176 | |
| 177 | def score(self, input_batch: NestedTensor) -> Tensor: |
| 178 | y = self.predict(input_batch) |
no test coverage detected