Skip to content

Commit a5c505d

Browse files
feat: Updated src/api.py
1 parent ac124aa commit a5c505d

File tree

1 file changed

+4
-4
lines changed

1 file changed

+4
-4
lines changed

src/api.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,11 @@
1-
from fastapi import FastAPI, UploadFile, File
2-
from PIL import Image
31
import torch
2+
from cnn import CNN # Importing CNN class from cnn.py
3+
from fastapi import FastAPI, File, UploadFile
4+
from PIL import Image
45
from torchvision import transforms
5-
from main import Net # Importing Net class from main.py
66

77
# Load the model
8-
model = Net()
8+
model = CNN()
99
model.load_state_dict(torch.load("mnist_model.pth"))
1010
model.eval()
1111

0 commit comments

Comments
 (0)