running on cpu

arguments.py

@@ -37,8 +37,6 @@ def get_args():
         '--infer_no_repeat_size', type=int, default=2, help="no repeat ngram size for inference")
     parser.add_argument(
         '--response-length', type=int, default=20, help='number of tokens to generate for each prompt.')
-    parser.add_argument(
-        '--num-gpus', type=int, default=None, help='number of gpus. use all available if none')
     parser.add_argument(
         '--port', type=int, default=None, help="port for the demo server")
 
@@ -47,11 +45,6 @@ def get_args():
 
     if args.use_label_prefix:
         log.info(f'using label prefix')
-    num_gpus = torch.cuda.device_count()
-    if args.num_gpus is None:
-        args.num_gpus = num_gpus
-    else:
-        args.num_gpus = min(num_gpus, args.num_gpus)
 
     if args.checkpoint is not None:
         args.checkpoint = str(Path(args.checkpoint).resolve())

load.py

@@ -42,7 +42,7 @@ def load_model(args, device, finetune=False):
                     use_transformer_mapper=args.use_transformer_mapper,
                     model_weight='None', use_label_prefix=args.use_label_prefix)
     ckpt = args.checkpoint + '.ckpt'
-    state = torch.load(ckpt)
+    state = torch.load(ckpt, map_location=torch.device('cpu'))
     policy_key = 'policy_model'
     if policy_key in state:
         policy.model.load_state_dict(state[policy_key])

run.py

@@ -22,16 +22,15 @@ log = logging.getLogger(__name__)
 
 
 def prepare(args):
-    num_gpus = torch.cuda.device_count()
-    log.info(f'Detect {num_gpus} GPUS')
     device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+    log.info(f'Device: {device}')
     args = load_model_args(args)
 
     def load_style(args, checkpoint):
         model = AutoModelForCausalLM.from_pretrained(args.init_model)
         if checkpoint is not None and Path(checkpoint).is_file():
             log.info("joint model: loading pretrained style generator")
-            state = torch.load(checkpoint)
+            state = torch.load(checkpoint, map_location=torch.device('cpu'))
             if 'global_step' in state:
                 step = state['global_step']
                 log.info(f'trained for {step} steps')