diff --git a/train_rsvqa.py b/train_rsvqa.py index 446105777eaf0f83edf5047851d337b9bd3de50a..6d1e448c71a2478859a9d5206fc9bb367cc4ccfb 100644 --- a/train_rsvqa.py +++ b/train_rsvqa.py @@ -269,6 +269,8 @@ def main( pl.seed_everything(seed, workers=True) + torch.set_float32_matmul_precision("medium") + img_size = 120 channels = 10 @@ -326,7 +328,6 @@ def main( logger=wandb_logger, check_val_every_n_epoch=2, callbacks=[checkpoint_callback, early_stopping_callback, lr_monitor], - precision="16", ) model = LitVisionEncoder(config=model_config, lr=lr)