| import torch |
| from BasePythonModelV3 import BasePythonModelV3 |
|
|
|
|
| class ECGRecoverPythonModelV3(BasePythonModelV3): |
| def predict(self, context, model_inputs: list[torch.Tensor]): |
| with torch.inference_mode(): |
| assert ( |
| isinstance(model_inputs, list) |
| and len(model_inputs) == 2 |
| and all([each_input.dim() == 3 for each_input in model_inputs]) |
| ), "expect list of 2 tensors as input: [digitized, mask], each tensor shape: (batch_size, n_leads, length)" |
| model_inputs = [ |
| each_input.to(self.device, self.dtype) for each_input in model_inputs |
| ] |
|
|
| self.network.eval() |
| logit = self.network(model_inputs) |
| output = self.network.activation(logit).cpu().numpy() |
| return output |
|
|
