Update model.py

model.py

@@ -1,12 +1,12 @@
+%%writefile /content/lung_disease_detection/model.py
 import torch
-import torchvision
 from torch import nn
 from torchvision import models
 
 class FineTunedResNet(nn.Module):
     def __init__(self, num_classes=3):
         super(FineTunedResNet, self).__init__()
-        self.resnet = models.resnet50(pretrained=True)
+        self.resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
         self.resnet.fc = nn.Sequential(
             nn.Linear(self.resnet.fc.in_features, 512),
             nn.ReLU(),
@@ -16,4 +16,3 @@ class FineTunedResNet(nn.Module):
 
     def forward(self, x):
         return self.resnet(x)
-