| 115 | @serve.ingress(app) |
| 116 | class MNISTClassifier: |
| 117 | def __init__(self): |
| 118 | self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| 119 | self.model = MNISTNet().to(self.device) |
| 120 | # Define the transformation pipeline for the input images. |
| 121 | self.transform = v2.Compose([ |
| 122 | v2.ToImage(), |
| 123 | v2.ToDtype(torch.float32, scale=True), |
| 124 | # Mean and standard deviation of the MNIST training subset. |
| 125 | v2.Normalize(mean=[0.1307], std=[0.3013]), |
| 126 | ]) |
| 127 | |
| 128 | self.model.eval() |
| 129 | |
| 130 | # batch_wait_timeout_s is the maximum time to wait for a full batch, |
| 131 | # trading off latency for throughput. |