Update app.py

app.py

@@ -1,5 +1,6 @@
 import torch
 from transformers import AutoModel
+import torch.nn as nn
 from PIL import Image
 import numpy as np
 import streamlit as st
@@ -30,6 +31,7 @@ if uploaded_file is not None:
         predictions = model(new_image)
         logits = predictions.last_hidden_state
         logits = logits.view(logits.shape[0], -1)
+        num_classes=2
         feature_reducer = nn.Linear(logits.shape[1], num_classes)
 
         logits = logits.to(device)