No module named 'upfirdn2d_plugin'

[hadaev8]

Describe the bug
Trying to use class AugmentPipe in my project.

To Reproduce
Run this notebook
https://colab.research.google.com/drive/1WyNx2jSNlXPBtXgpOmLQMxRE2GzFjeFM?usp=sharing

Trace

No module named 'upfirdn2d_plugin'
warnings.warn('Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:\n\n' + str(sys.exc_info()[1]))
Setting up PyTorch plugin "upfirdn2d_plugin"... Failed!
Setting up PyTorch plugin "upfirdn2d_plugin"... Failed!
Setting up PyTorch plugin "upfirdn2d_plugin"... Failed!
Setting up PyTorch plugin "upfirdn2d_plugin"... Failed!
tensor(0.8200, device='cuda:0', grad_fn=)
/content/ada/torch_utils/ops/upfirdn2d.py:34: UserWarning: Failed to build CUDA kernels for upfirdn2d. Falling back to slow reference implementation. Details:

Desktop (please complete the following information):
Google colab, linux, latest pytorch

[hadaev8]

#2 (comment)