Hi @ptrblck , thanks for your reply!
I have tried modifying my initial approach to:
structure = nn.ModuleList([nn.ModuleList([BranchRoutingModule(in_channels=self.in_channels) for j in range(int(pow(2, i)))]) for i in range(self.tree_height - 1)]) 

I still get the same errror: "“Input type (torch.cuda.FloatTensor) and weight type (torch.FloatTensor) should be the same”…
Running out of ideas here…
Just for extra precision, here is my branchrouting module:
class BranchRoutingModule(nn.Module):

  def __init__(self, in_channels=512, epsilon=1e-12):
    super(BranchRoutingModule, self).__init__()

    self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1)
    self.l2norm = lambda x: torch.nn.functional.normalize(x, dim=1)
    self.fc = nn.Linear(in_channels, 1)
    self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) # size 1x1
    self.signedsqrt = lambda x: torch.sign(x) * torch.sqrt(torch.sign(x) * x + epsilon) 
    self.sigmoid = nn.Sigmoid()

  def forward(self, x):
    print("BR X : " + str(x.is_cuda))
    feature = self.conv1(x)
    feature = self.avgpool(feature)
    feature = self.signedsqrt(feature)
    feature = self.l2norm(feature)
    feature = feature.view(feature.size(0), -1)
    out = self.fc(feature)
    out = self.sigmoid(out)

    return out

Thanks!