Update Apollo/inference.py

Apollo/inference.py

@@ -10,6 +10,7 @@ from ml_collections import ConfigDict
 import json
 import time
 import warnings
+import spaces
 warnings.filterwarnings("ignore")
 
 def get_config(config_path):
@@ -124,6 +125,7 @@ def main(input_wav, output_wav, ckpt_path):
     del model
     torch.cuda.empty_cache()
 
+@spaces.GPU
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Audio Inference Script")
     parser.add_argument("--in_wav", type=str, required=True, help="Path to input wav file")