The counting of 'NFE-F' does not work in the case of multiple GPUs

[HanshuYAN]

Hi, I simply modified the code, odenet_mnist.py, by adding 'model = nn.DataParallel(model, devices)'.

Then, the code can be run on multiple GPUs, but the counting of NFE-F kept unchanged, always 0. Do you know why?

[rtqichen]

This is because the replicated models don't support mutation. See the first warning in https://pytorch.org/docs/stable/nn.html?highlight=nn%20dataparallel#dataparallel-layers-multi-gpu-distributed.

Changing this line https://github.com/rtqichen/torchdiffeq/blob/master/examples/odenet_mnist.py#L102 to self.register_buffer("nfe", torch.tensor(0.)) will make it work with nn.DataParallel, but it'll only count the number of evaluations on GPU 0.

[HanshuYAN]

Thx~