When running segmentation predict on gpu, the conversion from tensor to numpy fails. Calling `.cpu()` solves this problem. 
Signed-off-by: Paul Guerrie <97041392+paulguerrie@users.noreply.github.com>
Signed-off-by: Paul Guerrie <97041392+paulguerrie@users.noreply.github.com>