remove to(device) calls in load_model

inference.py

@@ -161,10 +161,6 @@ def load_model(args):
     # as these weights are only used for inference, keeping weights in full precision is not required.
     weight_dtype = torch.bfloat16
 
-    text_encoder.to(args.device, dtype=weight_dtype)
-    transformer.to(args.device, dtype=weight_dtype)
-    vae.to(args.device, dtype=weight_dtype)
-
     pipe = ControlnetCogVideoXPipeline.from_pretrained(
         args.pretrained_model_path,
         tokenizer=tokenizer,
@@ -186,7 +182,6 @@ def load_model(args):
         scheduler_args["variance_type"] = variance_type
 
     pipe.scheduler = CogVideoXDPMScheduler.from_config(pipe.scheduler.config, **scheduler_args)
-    pipe = pipe.to(args.device)
 
     return pipe, model_config
 
@@ -247,6 +242,11 @@ def main(args):
 
     pipe, model_config = load_model(args)
 
+    # text_encoder.to(args.device, dtype=weight_dtype)
+    # transformer.to(args.device, dtype=weight_dtype)
+    # vae.to(args.device, dtype=weight_dtype)    
+    pipe = pipe.to(args.device)
+
     for image_path in image_paths:
         image = Image.open(image_path)