can u update this for the new pytorch
can u update this for the new pytorch