Pytorch model wrapper