Kohaku-Blueleaf commited on
Commit
d71fa96
·
1 Parent(s): 4393fe6
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -2,6 +2,7 @@ import os
2
  import random
3
  import json
4
  from pathlib import Path
 
5
 
6
  if os.environ.get("IN_SPACES", None) is not None:
7
  in_spaces = True
@@ -60,7 +61,9 @@ DEFAULT_FORMAT = """
60
  """.strip()
61
 
62
 
63
- def GPU(func, duration=None):
 
 
64
  if in_spaces:
65
  if duration:
66
  return spaces.GPU(func, duration=duration)
 
2
  import random
3
  import json
4
  from pathlib import Path
5
+ from functools import partial
6
 
7
  if os.environ.get("IN_SPACES", None) is not None:
8
  in_spaces = True
 
61
  """.strip()
62
 
63
 
64
+ def GPU(func=None, duration=None):
65
+ if func is None:
66
+ return partial(GPU, duration=duration)
67
  if in_spaces:
68
  if duration:
69
  return spaces.GPU(func, duration=duration)