@Chen0729
I guess the problem is here:
bn1_b = register_parameter(“bn1_b”,torch::zeros(500));
bn2_b = register_parameter(“bn2_b”,torch::zeros(500));
Both zero tensor are CPU tensors.
x was converted to gpu but bn1_b is not.
x=torch::nn::functional::softplus(bn1(fc1(x))+bn1_b)
And I am wondering what is the error message when you do assert? assert should be doable.