Where to specify the device?
marian42 opened this issue · comments
From the readme:
You can also filter the device to report on by passing extra arguments:
report(device=torch.device(0))
The device parameter seems to work neither with the report decorator nor with the MemReporter. If this feature is still available, the readme should be clarified.
The device= filter is added in the latest master branch, you can install via:
pip install git+https://github.com/stonesjtu/pytorch_memlab
I just uploaded the current master to PyPI. you can try to update the package by pip install -U pytorch_memlab after several minutes.
Thank you!