Where is the nn.Linear cuda implementation

Hi @Ziyu_Huang , I had the same issue so I wrote code to implement the Conv2D and Linear layers using native CUDA code here. Hopefully it will be useful for you! I only wrote code for the forward pass in CUDA but hopefully you can use that as a base to write code for backward if you need that.

1 Like