The error points to a size mismatch in dim2 for block2 and block1.
I just skimmed through the code, but maybe you wanted to use block1_pool instead of block1?
Generally, all shapes besides the specified dimension should match, when you are trying to use torch.cat.
In your case, dim1 could differ, but all other dims should have the same shape.