Spaces:
Runtime error
Runtime error
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
|