Spaces:

fun-research
/

TiTok

Running on Zero

yucornetto commited on Jun 21

Commit

506552e

•

1 Parent(s): 062a2e2

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -11,7 +11,7 @@ import argparse
 import demo_util
 import os
-device = "cpu"
 model2ckpt = {
     "TiTok-L-32": ("tokenizer_titok_l32.bin", "generator_titok_l32.bin"),

 import demo_util
 import os
+device = "cuda" if torch.cuda.is_available() else "cpu"
 model2ckpt = {
     "TiTok-L-32": ("tokenizer_titok_l32.bin", "generator_titok_l32.bin"),