Update app.py

app.py

@@ -5,6 +5,7 @@ from transformers import DPTImageProcessor
 import gradio as gr
 import matplotlib.pyplot as plt
 from mpl_toolkits.mplot3d import Axes3D
+import torch.nn as nn
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")