Move wandb import

Author: palonso

src/modules/maskingmodel.py

@@ -6,7 +6,6 @@
 
 import gin
 import torch
-import wandb
 from torch import nn
 import pytorch_lightning as L
 
@@ -407,6 +406,8 @@ def training_step(self, batch, batch_idx):
                 )
 
         elif self.first_coverage and batch_idx == first_coverage_steps:
+            import wandb
+
             # Print the histogram you can check it in the wandb dashboard (log section)
             for key, value in self.tokens_accumulator.items():
                 self.logger.experiment.log({f"{key}_histogram": wandb.Histogram(value)})