Replace FFNN in GCN-FFNN model with other regressor

Hi community, thanks for the wonderful work that you are doing. Is it possible to replace the nn.Linear in the code below with other common regressors like RandomForest from Sckit-Learn? Any suggestion on how?

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import dgl
from dgl.nn import GraphConv
from sklearn.ensemble import RandomForestRegressor

class GCN(nn.Module):
    def __init__(self, in_feats, h_feats, num_classes):
        super(GCN, self).__init__()
        if  torch.cuda.is_available():
            self.device = torch.device('cuda:0')
            self.device = torch.device('cpu')
        self.conv = GraphConv(in_feats, h_feats)
        if torch.cuda.is_available():
            self.conv = self.conv.cuda()
        self.readout1 = nn.Linear(h_feats, h_feats).to(self.device)
        self.readout2 = nn.Linear(h_feats, num_classes).to(self.device)

    def forward(self, g, in_feat):
        with g.local_scope():
            h = self.conv(g, in_feat)
            h = F.relu(h)
            g.ndata['h'] = h
            meanNodes = dgl.mean_nodes(g, 'h')
            output = F.relu(self.readout1(meanNodes))
            output = self.readout2(output)
        return output

Replacing it with sklearn.RandomForestRegressor will break auto-differentiation, so the model will not train properly. Maybe this post on PyTorch forum will be helpful? How can I use KNN, Random Forest models in Pytorch? - #3 by pydv - vision - PyTorch Forums