Just installed pytorch-nightly in a dedicated conda env for GPU and then CPU only (since the first option didn’t work, I suspected a downgrade imposed by my cuda 10.1) in order to use torch.vmap but it seems I don’t have the right nightly: I keep generating the error

AttributeError: module ‘torch’ has no attribute ‘vmap’

Furthermore, a
import torch
print(torch.__version__)

in the notebook launched with the new nightly dedicated conda env returns 1.7.1, although the nightly should be 1.8.X according to this post. I tried installing pytorch nightly successively with the command lines suggested on the pytorch website homepage:
conda install pytorch torchvision torchaudio cudatoolkit=10.1 -c pytorch-nightly
conda install pytorch torchvision torchaudio cpuonly -c pytorch-nightly
Probably unrelated but would it be linked to the Remove torch.vmap here 1 ? According to this github issue torch.vmap should be available in pytorch 1.7 anyway.
I’m calling torch.vmap in the following line of code, part of a class method:
torch.vmap(self.distance_riemann_to_c)(net_outputs)
torch.vmap thus vectorizes the application of the other method (same class) self.distance_riemann_to_c to the batched nnet outputs net_outputs. Sorry if it would have been more appropriate to open a new post (?).