RuntimeError: size mismatch, m1: [1 x 588], m2: [12 x 10] at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:268 about flashtorch HOT 5 CLOSED

def conv_block(in_channels, out_channels, k):
    return nn.Sequential(
        nn.Conv2d(in_channels, in_channels, k, padding=0),
        nn.BatchNorm2d(in_channels),
        nn.ReLU(),
        nn.MaxPool2d(2)
    )

class Top(nn.Module):
  def __init__(self):
    super().__init__()
    self.encoder = conv_block(3, 3, 1)
    self.lin = nn.Linear(20, 10)
    self.childone = Second()
    self.childtwo = Second()
  def forward(self, x):
    # set_trace()
    a, b = self.childone(self.encoder(x)), self.childtwo(self.encoder(x))
    # print('top', a.shape, b.shape)
    out = torch.cat((a, b), dim=-1)
    return self.lin(out) 

class Second(nn.Module):
  def __init__(self):
    super().__init__()
    self.encoder = conv_block(3, 3, 1)
    self.lin = nn.Linear(20, 10)
    self.childone = Middle()
    self.childtwo = Middle()

  def forward(self, x):
    a, b = self.childone(self.encoder(x)), self.childtwo(self.encoder(x))
    # print('middle', a.shape, b.shape)
    out = torch.cat((a, b), dim=-1)
    return self.lin(out)

class Middle(nn.Module):
  def __init__(self):
    super().__init__()
    self.encoder = conv_block(3, 3, 1)
    self.lin = nn.Linear(20, 10)
    self.childone = Bottom()
    self.childtwo = Bottom()

  def forward(self, x):
    a, b = self.childone(self.encoder(x)), self.childtwo(self.encoder(x))
    # print('middle', a.shape, b.shape)
    out = torch.cat((a, b), dim=-1)
    return self.lin(out)

class Bottom(nn.Module):
  def __init__(self):
    super().__init__()
    self.encoder = conv_block(3, 3, 1)
    self.lin_one = nn.Linear(12, 10)
  def forward(self, x):
    # print('bottom', x.shape)
    out = self.encoder(x)
    return (self.lin_one(out.view(out.size(0), -1)))

model = Top()
model.to('cuda')
from flashtorch.activmax import GradientAscent
g_ascent = GradientAscent(model)
g_ascent.use_gpu = True

g_ascent.visualize(model.childtwo.childtwo.childtwo.encoder[0], title='conv');

g_ascent = GradientAscent(model, img_size=int)

g_ascent.img_size = int