@@ -104,7 +104,10 @@ def extract_body_pose(in_folder, is_video, image_shape, column):
model = get_model('vgg19')
model.load_state_dict(torch.load(weight_name))
- model = torch.nn.DataParallel(model).cuda()
+ if torch.cuda.is_available():
+ model = torch.nn.DataParallel(model).cuda()
+ else:
+ model = torch.nn.DataParallel(model)
model.float()
model.eval()
@@ -153,7 +156,8 @@ def extract_body_id(in_folder, is_video, image_shape, column):
net.load_state_dict(torch.load('reid_baseline/reid_model.pth'))
net.model.fc = nn.Sequential()
net.classifier = nn.Sequential()
- net.cuda()
+ if torch.cuda.is_available():
+ net.cuda()
data_transforms = transforms.Compose([
transforms.ToPILImage(),
--
2.39.0.windows.2