Made loggers togglable from top level

symbolic_nn_tests/__main__.py

@@ -1,25 +1,36 @@
 LEARNING_RATE = 10e-5
 
 
-def qmnist_test(loss_func, version):
+def qmnist_test(loss_func, version, tensorboard=True, wandb=True):
     from .experiment_1.model import main as test_model
+
+    logger = []
+
+    if tensorboard:
         from lightning.pytorch.loggers import TensorBoardLogger
-    from lightning.pytorch.loggers import WandbLogger
-    import wandb
 
         tb_logger = TensorBoardLogger(
             save_dir=".",
             name="logs/comparison",
             version=version,
         )
+        logger.append(tb_logger)
+
+    if wandb:
+        import wandb as _wandb
+        from lightning.pytorch.loggers import WandbLogger
+
         wandb_logger = WandbLogger(
             project="Symbolic_NN_Tests",
             name=version,
             dir="wandb",
         )
-    logger = [tb_logger, wandb_logger]
+        logger.append(wandb_logger)
+
     test_model(logger=logger, loss_func=loss_func, lr=LEARNING_RATE)
-    wandb.finish()
+
+    if wandb:
+        _wandb.finish()
 
 
 def qmnist_experiment():