patch

diff --git a/models/hierarchy_inference_model.py b/models/hierarchy_inference_model.py
index 3116307..5de661d 100644
--- a/models/hierarchy_inference_model.py
+++ b/models/hierarchy_inference_model.py
@@ -21,7 +21,7 @@ class VQGANTextureAwareSpatialHierarchyInferenceModel():
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.is_train = opt['is_train']
 
         self.top_encoder = Encoder(
diff --git a/models/hierarchy_vqgan_model.py b/models/hierarchy_vqgan_model.py
index 4b0d657..0bf4712 100644
--- a/models/hierarchy_vqgan_model.py
+++ b/models/hierarchy_vqgan_model.py
@@ -20,7 +20,7 @@ class HierarchyVQSpatialTextureAwareModel():
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.top_encoder = Encoder(
             ch=opt['top_ch'],
             num_res_blocks=opt['top_num_res_blocks'],
diff --git a/models/parsing_gen_model.py b/models/parsing_gen_model.py
index 9440345..15a1ecb 100644
--- a/models/parsing_gen_model.py
+++ b/models/parsing_gen_model.py
@@ -22,7 +22,7 @@ class ParsingGenModel():
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.is_train = opt['is_train']
 
         self.attr_embedder = ShapeAttrEmbedding(
diff --git a/models/sample_model.py b/models/sample_model.py
index 4c60e3f..5265cd0 100644
--- a/models/sample_model.py
+++ b/models/sample_model.py
@@ -23,7 +23,7 @@ class BaseSampleModel():
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
 
         # hierarchical VQVAE
         self.decoder = Decoder(
@@ -123,7 +123,7 @@ class BaseSampleModel():
 
     def load_top_pretrain_models(self):
         # load pretrained vqgan
-        top_vae_checkpoint = torch.load(self.opt['top_vae_path'])
+        top_vae_checkpoint = torch.load(self.opt['top_vae_path'], map_location=self.device)
 
         self.decoder.load_state_dict(
             top_vae_checkpoint['decoder'], strict=True)
@@ -137,7 +137,7 @@ class BaseSampleModel():
         self.top_post_quant_conv.eval()
 
     def load_bot_pretrain_network(self):
-        checkpoint = torch.load(self.opt['bot_vae_path'])
+        checkpoint = torch.load(self.opt['bot_vae_path'], map_location=self.device)
         self.bot_decoder_res.load_state_dict(
             checkpoint['bot_decoder_res'], strict=True)
         self.decoder.load_state_dict(checkpoint['decoder'], strict=True)
@@ -153,7 +153,7 @@ class BaseSampleModel():
 
     def load_pretrained_segm_token(self):
         # load pretrained vqgan for segmentation mask
-        segm_token_checkpoint = torch.load(self.opt['segm_token_path'])
+        segm_token_checkpoint = torch.load(self.opt['segm_token_path'], map_location=self.device)
         self.segm_encoder.load_state_dict(
             segm_token_checkpoint['encoder'], strict=True)
         self.segm_quantizer.load_state_dict(
@@ -166,7 +166,7 @@ class BaseSampleModel():
         self.segm_quant_conv.eval()
 
     def load_index_pred_network(self):
-        checkpoint = torch.load(self.opt['pretrained_index_network'])
+        checkpoint = torch.load(self.opt['pretrained_index_network'], map_location=self.device)
         self.index_pred_guidance_encoder.load_state_dict(
             checkpoint['guidance_encoder'], strict=True)
         self.index_pred_decoder.load_state_dict(
@@ -176,7 +176,7 @@ class BaseSampleModel():
         self.index_pred_decoder.eval()
 
     def load_sampler_pretrained_network(self):
-        checkpoint = torch.load(self.opt['pretrained_sampler'])
+        checkpoint = torch.load(self.opt['pretrained_sampler'], map_location=self.device)
         self.sampler_fn.load_state_dict(checkpoint, strict=True)
         self.sampler_fn.eval()
 
@@ -397,7 +397,7 @@ class SampleFromPoseModel(BaseSampleModel):
                         [185, 210, 205], [130, 165, 180], [225, 141, 151]]
 
     def load_shape_generation_models(self):
-        checkpoint = torch.load(self.opt['pretrained_parsing_gen'])
+        checkpoint = torch.load(self.opt['pretrained_parsing_gen'], map_location=self.device)
 
         self.shape_attr_embedder.load_state_dict(
             checkpoint['embedder'], strict=True)
diff --git a/models/transformer_model.py b/models/transformer_model.py
index 7db0f3e..4523d17 100644
--- a/models/transformer_model.py
+++ b/models/transformer_model.py
@@ -21,7 +21,7 @@ class TransformerTextureAwareModel():
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.is_train = opt['is_train']
 
         # VQVAE for image
@@ -317,10 +317,10 @@ class TransformerTextureAwareModel():
     def sample_fn(self, temp=1.0, sample_steps=None):
         self._denoise_fn.eval()
 
-        b, device = self.image.size(0), 'cuda'
+        b = self.image.size(0)
         x_t = torch.ones(
-            (b, np.prod(self.shape)), device=device).long() * self.mask_id
-        unmasked = torch.zeros_like(x_t, device=device).bool()
+            (b, np.prod(self.shape)), device=self.device).long() * self.mask_id
+        unmasked = torch.zeros_like(x_t, device=self.device).bool()
         sample_steps = list(range(1, sample_steps + 1))
 
         texture_mask_flatten = self.texture_tokens.view(-1)
@@ -336,11 +336,11 @@ class TransformerTextureAwareModel():
 
         for t in reversed(sample_steps):
             print(f'Sample timestep {t:4d}', end='\r')
-            t = torch.full((b, ), t, device=device, dtype=torch.long)
+            t = torch.full((b, ), t, device=self.device, dtype=torch.long)
 
             # where to unmask
             changes = torch.rand(
-                x_t.shape, device=device) < 1 / t.float().unsqueeze(-1)
+                x_t.shape, device=self.device) < 1 / t.float().unsqueeze(-1)
             # don't unmask somewhere already unmasked
             changes = torch.bitwise_xor(changes,
                                         torch.bitwise_and(changes, unmasked))
diff --git a/models/vqgan_model.py b/models/vqgan_model.py
index 13a2e70..9c840f1 100644
--- a/models/vqgan_model.py
+++ b/models/vqgan_model.py
@@ -20,7 +20,7 @@ class VQModel():
     def __init__(self, opt):
         super().__init__()
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.encoder = Encoder(
             ch=opt['ch'],
             num_res_blocks=opt['num_res_blocks'],
@@ -390,7 +390,7 @@ class VQImageSegmTextureModel(VQImageModel):
 
     def __init__(self, opt):
         self.opt = opt
-        self.device = torch.device('cuda')
+        self.device = torch.device(opt['device'])
         self.encoder = Encoder(
             ch=opt['ch'],
             num_res_blocks=opt['num_res_blocks'],