How to initialize the parameter of nn.LazyLinear since I have to use dataparallel to push the model onto?

Hi there,

I’m currently using a LazyLinear layer in my network. I need to push the network to 2 GPUs, so I need to use DataParallel (model).
It seems like I cannot initialize the LazyLinear layer before I input dummy data into the model. But if I:
input dummy data
→ use dataparallel
→ input real data to train the model,
Another error occurred.

RuntimeError: CUDA error: CUBLAS_STATUS_EXECUTION_FAILED when calling `cublasSgemm( handle, opa, opb, m, n, k, &alpha, a, lda, b, ldb, &beta, c, ldc)`

Could anyone tell me how to use both lazylinear and dataparallel?

Thank you!