Update main.py

main.py

@@ -1,4 +1,5 @@
 import os
+import spaces
 import torch
 import numpy as np
 import argparse
@@ -9,6 +10,7 @@ from utils import load_image, load_mask, load_mask_edit
 from utils_mask import process_mask_move_torch, process_mask_remove_torch, mask_union_torch, mask_substract_torch, create_outer_edge_mask_torch
 from utils_mask import check_mask_overlap_torch, check_cover_all_torch, visualize_mask_list, get_mask_difference_torch, save_mask_list_to_npys
 
+@spaces.GPU
     
 def run_main(
     name="example_tmp",