add the torch import in model.py

model.py

@@ -1,3 +1,4 @@
+import torch
 import torchvision
 from torchvision import transforms
 from PIL import Image
@@ -45,9 +46,9 @@ def create_model(num_classes: int = 20,
   torch.manual_seed(seed)
 
   # 6. Modify the number of output layers
-  model.heads = nn.Sequential(
-      nn.Dropout(p=0.5),
-      nn.Linear(in_features=768, out_features=num_classes, bias=True)
+  model.heads = torch.nn.Sequential(
+      torch.nn.Dropout(p=0.5),
+      torch.nn.Linear(in_features=768, out_features=num_classes, bias=True)
   )
 
   return model, combined_transforms