RuntimeError: mat1 and mat2 shapes cannot be multiplied (GaussianParametrizer)

[dineshdaultani]

Hi @anujinho,
I am trying to reproduce your TRIDENT CCVAE model, however, I am not able to pass inputs through the learner/model.
Below is the model:

MAML(
  (module): CCVAE(
    (encoder): CEncoder(
      (net): Sequential(
        (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (2): LeakyReLU(negative_slope=0.2)
        (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (5): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (6): LeakyReLU(negative_slope=0.2)
        (7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (8): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (9): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (10): LeakyReLU(negative_slope=0.2)
        (11): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (12): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
        (14): LeakyReLU(negative_slope=0.2)
        (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (16): Flatten(start_dim=1, end_dim=-1)
      )
    )
    (decoder): CDecoder(
      (linear): Sequential(
        (0): Linear(in_features=128, out_features=800, bias=True)
        (1): LeakyReLU(negative_slope=0.2)
      )
      (net): Sequential(
        (0): UpsamplingNearest2d(size=(10, 10), mode=nearest)
        (1): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
        (2): LeakyReLU(negative_slope=0.2)
        (3): UpsamplingNearest2d(size=(21, 21), mode=nearest)
        (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
        (5): LeakyReLU(negative_slope=0.2)
        (6): UpsamplingNearest2d(size=(42, 42), mode=nearest)
        (7): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=same)
        (8): LeakyReLU(negative_slope=0.2)
        (9): UpsamplingNearest2d(size=(84, 84), mode=nearest)
        (10): Conv2d(32, 3, kernel_size=(3, 3), stride=(1, 1), padding=same)
        (11): Sigmoid()
      )
    )
    (classifier_vae): Classifier_VAE(
      (encoder): TADCEncoder(
        (net): Sequential(
          (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (2): LeakyReLU(negative_slope=0.2)
          (3): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (5): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (6): LeakyReLU(negative_slope=0.2)
          (7): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (8): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (9): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (10): LeakyReLU(negative_slope=0.2)
          (11): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
          (12): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
          (13): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
          (14): LeakyReLU(negative_slope=0.2)
          (15): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        )
        (fe): Sequential(
          (0): Conv2d(1, 64, kernel_size=(110, 1), stride=(1, 1), padding=valid, bias=False)
          (1): LeakyReLU(negative_slope=0.2)
          (2): Conv2d(64, 32, kernel_size=(1, 1), stride=(1, 1), padding=valid, bias=False)
          (3): LeakyReLU(negative_slope=0.2)
        )
        (f_q): Conv2d(32, 1, kernel_size=(1, 1), stride=(1, 1), padding=valid, bias=False)
        (f_k): Conv2d(32, 1, kernel_size=(1, 1), stride=(1, 1), padding=valid, bias=False)
        (f_v): Conv2d(32, 1, kernel_size=(1, 1), stride=(1, 1), padding=valid, bias=False)
      )
      (gaussian_parametrizer): GaussianParametrizer(
        (h1): Linear(in_features=864, out_features=64, bias=True)
        (h2): Linear(in_features=864, out_features=64, bias=True)
      )
      (classifier): Sequential(
        (0): Linear(in_features=64, out_features=32, bias=True)
        (1): LeakyReLU(negative_slope=0.2)
        (2): Linear(in_features=32, out_features=10, bias=True)
      )
    )
    (gaussian_parametrizer): GaussianParametrizer(
      (h1): Linear(in_features=800, out_features=64, bias=True)
      (h2): Linear(in_features=800, out_features=64, bias=True)
    )
  )
)

And below is the minimized stack trace:

Traceback (most recent call last):
  ...
  File ".../trident.py", line 105, in _train_epoch
    eval_loss, eval_acc = inner_adapt_trident(ttask, reconst_loss, 
  File ".../utils.py", line 146, in inner_adapt_trident
    reconst_image, logits, mu_l, log_var_l, mu_s, log_var_s = learner(
 ...
  File ".../archs.py", line 812, in forward
    mu_s, log_var_s = self.gaussian_parametrizer(xs)
  ...
  File ".../archs.py", line 404, in forward
    mu = self.h1(x)
  ...
RuntimeError: mat1 and mat2 shapes cannot be multiplied (10x128 and 800x64)

Also, below are the hyperparameters:

n_ways: 5
k_shots: 1 
q_shots: 10 
meta_batch_size: 20
order: False
inner_lr: 0.0014
task_adapt: True
zl: 64
zs: 64
reconstr: std
dataset: cifarfs
wm_channels: 64
wn_channels: 32
download: False
extra: False
adapt_steps_train: 5
adapt_steps_test: 5

Basically the output shape of CEncoder and input shape of GaussianParametrizer doesn't match.
Can you please help to resolve this issue. Will look forward to hearing from you soon. Thanks!