I am writing a custom linear layer and I want to execute it on GPU. The tensors are correctly located on the GPU ( *.is_cuda returns true for input & weight). I am performing multiplications directly using two loops and I am getting the same answer as the in-built PyTorch Linear layer, but I’m not sure if this calculation is actually being done on the GPU. Is there a way to check?

def forward(ctx, X, weight, bias):
    ctx.save_for_backward(X, weight, bias)

    print(f"Input:: GPU: {X.is_cuda}, Weight:: GPU: {weight.is_cuda})

    (m, n_W) = X.shape
    (A, _) = weight.shape

    # Transpose the weights
    weight = torch.transpose(weight, 0, 1)

    output = torch.cuda.FloatTensor(m, A).fill_(0)

    for k in range(m):
        for l in range(A):
            accumulation = 0.0
            for j in range(n_W):
                accumulation += X[k, j] * weight[j, l]
            output[k, l] = accumulation + bias[l]

    return output


Thanks!