Is it possible to send a model to a cuda kernel?

1

Hello,

I am new with pytorch and I am trying to implement pytorch c++ with a cuda/c++ project.

My question is the following, is it possible to do something like this:

//// Neural Net. definition ////

struct Net : torch::nn::Module
{
Net(int inputSize, int hiddenSize, int numClasses):
fc1(register_module(“fc1”, torch::nn::Linear(inputSize, hiddenSize))),
fc2(register_module(“fc2”, torch::nn::Linear(hiddenSize, numClasses))) {};

torch::nn::Linear fc1;
torch::nn::Linear fc2;

torch::Tensor forward(torch::Tensor input) 
{
    auto out = fc1(input);
    out = torch::relu(out);
    out = fc2(out);
    
    return out;    

}

};

//// Use Neural Net in CUDA kernel

global void testKernel (Net model)
{
… use model here …
}

Thanks for any help!!

1 Like