Update app.py

app.py

@@ -28,7 +28,7 @@ delta_ckpt = torch.load(args['delta_ckpt_path'], map_location=torch.device('cpu'
 model.load_state_dict(delta_ckpt, strict=False)
 delta_ckpt = torch.load(args['anomalygpt_ckpt_path'], map_location=torch.device('cpu'))
 model.load_state_dict(delta_ckpt, strict=False)
-model = model.eval()
+model = model.eval().to(torch.bfloat16)
 
 
 output = None
@@ -126,7 +126,7 @@ def predict(
     history.append((input, response))
 
 
-    plt.imshow(pixel_output.reshape(224,224).detach().cpu(), cmap='binary_r')
+    plt.imshow(pixel_output.to(torch.float16).reshape(224,224).detach().cpu(), cmap='binary_r')
     plt.axis('off')
     plt.savefig('output.png',bbox_inches='tight',pad_inches = 0)