Update app.py

app.py

@@ -8,7 +8,7 @@ import gradio as gr
 class MyModel(torch.nn.Module):
     def __init__(self, num_classes):
         super(MyModel, self).__init__()
-        self.model = models.resnet18(pretrained=False)  # ResNet18 모델 사용
+        self.model = models.resnet18(weights=None)  # ResNet18 모델 사용
         self.model.fc = torch.nn.Linear(self.model.fc.in_features, num_classes)  # 마지막 레이어 수정
 
     def forward(self, x):
@@ -17,8 +17,9 @@ class MyModel(torch.nn.Module):
 # 모델 초기화
 num_classes = 6  # 클래스 수에 맞게 설정
 model = MyModel(num_classes)
-# 모델 가중치 로드
-model.load_state_dict(torch.load('resnet18_finetuned_2.pth', weights_only=True), strict=False)
+
+# 모델 가중치 로드 (CPU로 매핑)
+model.load_state_dict(torch.load('resnet18_finetuned_2.pth', map_location=torch.device('cpu')), strict=False)
 model.eval()
 
 # 데이터 변환 정의