Spaces:
Running
on
Zero
Running
on
Zero
yucornetto
commited on
Commit
•
062a2e2
1
Parent(s):
3469bba
Update demo_util.py
Browse files- demo_util.py +2 -2
demo_util.py
CHANGED
@@ -37,14 +37,14 @@ def get_config(config_path):
|
|
37 |
|
38 |
def get_titok_tokenizer(config):
|
39 |
tokenizer = TiTok(config)
|
40 |
-
tokenizer.load_state_dict(torch.load(config.experiment.tokenizer_checkpoint,
|
41 |
tokenizer.eval()
|
42 |
tokenizer.requires_grad_(False)
|
43 |
return tokenizer
|
44 |
|
45 |
def get_titok_generator(config):
|
46 |
generator = ImageBert(config)
|
47 |
-
generator.load_state_dict(torch.load(config.experiment.generator_checkpoint,
|
48 |
generator.eval()
|
49 |
generator.requires_grad_(False)
|
50 |
return generator
|
|
|
37 |
|
38 |
def get_titok_tokenizer(config):
|
39 |
tokenizer = TiTok(config)
|
40 |
+
tokenizer.load_state_dict(torch.load(config.experiment.tokenizer_checkpoint, map_location="cpu"))
|
41 |
tokenizer.eval()
|
42 |
tokenizer.requires_grad_(False)
|
43 |
return tokenizer
|
44 |
|
45 |
def get_titok_generator(config):
|
46 |
generator = ImageBert(config)
|
47 |
+
generator.load_state_dict(torch.load(config.experiment.generator_checkpoint, map_location="cpu"))
|
48 |
generator.eval()
|
49 |
generator.requires_grad_(False)
|
50 |
return generator
|