RuntimeError: size mismatch, m1: [512 x 4], m2: [2048 x 512]

nn.Flatten() is part of PyTorch 1.2.0 and later. Either update your pytorch version or you may have to find alternate ways of doing that. More on that here Flatten layer of PyTorch build by sequential container