How to use inference on GPU? feature, _ = model(bgr_tensor_input) don't use GPU

[martinenkoEduard]

How to use inference on GPU? feature, _ = model(bgr_tensor_input) don't use GPU

[macqueen09]

after train your model:

def load_pretrained_model():
model = net.build_model("ir_101")
statedict = torch.load(MODEL_WEIGHT, map_location=device)['state_dict']
model_statedict = {key[6:]: val for key, val in statedict.items() if key.startswith('model.')}
model.load_state_dict(model_statedict)
model.eval()
return model.to(device)

model = load_pretrained_model()
model.eval()

feature, _ = model(batch)

(batch already in gpu )