How to optimize CPU usage by modifying test_EDTalk_A_use_exp_weight function to leverage GPU

[emboss33]

I noticed a significant increase in CPU usage in the following part of demo_EDTalk_A_using_predefined_exp_weights.py

I suspect that the function test_EDTalk_A_use_exp_weight in generator.py is causing this spike in CPU usage. I am planning to modify this function to utilize GPU (CUDA) more efficiently in order to reduce the load on the CPU and improve the calculation speed, anyone here?

img_target_lip = self.lip_vid_target[i:i+1]
if i >= len_pose:
    img_target_pose = self.pose_vid_target[:, -1, :, :, :]
else:
    img_target_pose = self.pose_vid_target[:, i, :, :, :]

img_recon = self.gen.test_EDTalk_A_use_exp_weight(self.img_source, img_target_lip, img_target_pose, self.exp_vid_target, h_start)

vid_target_recon.append(img_recon.unsqueeze(2))

[tanshuai0219]

I noticed a significant increase in CPU usage in the following part of demo_EDTalk_A_using_predefined_exp_weights.py

I suspect that the function test_EDTalk_A_use_exp_weight in generator.py is causing this spike in CPU usage. I am planning to modify this function to utilize GPU (CUDA) more efficiently in order to reduce the load on the CPU and improve the calculation speed, anyone here?

img_target_lip = self.lip_vid_target[i:i+1]
if i >= len_pose:
    img_target_pose = self.pose_vid_target[:, -1, :, :, :]
else:
    img_target_pose = self.pose_vid_target[:, i, :, :, :]

img_recon = self.gen.test_EDTalk_A_use_exp_weight(self.img_source, img_target_lip, img_target_pose, self.exp_vid_target, h_start)

vid_target_recon.append(img_recon.unsqueeze(2))

I have checked the codes, and the current models and function calculation are on GPU. I think only loading data depends on CPU.