hybrid encoder hasn't attn_layers

Author: TITC

pix2tex/models/utils.py

@@ -46,7 +46,9 @@ def get_model(args, training=False):
     decoder.to(args.device)
     if args.wandb:
         import wandb
-        en_attn_layers = encoder.module.attn_layers if available_gpus > 1 else encoder.attn_layers
+        en_attn_layers = encoder
+        if args.encoder_structure.lower() == 'vit':
+            en_attn_layers = encoder.module.attn_layers if available_gpus > 1 else encoder.attn_layers
         de_attn_layers = decoder.module.net.attn_layers if available_gpus > 1 else decoder.net.attn_layers
         wandb.watch((en_attn_layers, de_attn_layers))
     model = Model(encoder, decoder, args)