fix(research): remove wandb dependency and guard optional imports

- Strip wandb imports and calls from MultiProxAn diffusion models and
utils (prevents ImportError at checkpoint load time)
- Wrap graph_tool, pyemd, pygsp, dist_helper imports in try/except
across spectre_utils, molecular_metrics, train_metrics (these are
only needed for training metrics, not inference)
- Fix pandas deprecation in COINs load_graph.py (iteritems -> items,
to_frame column naming)

src/research/COINs-KGGeneration/graph_completion/graphs/load_graph.py

@@ -128,8 +128,8 @@ class Loader:
         hits_at_3_limit = community_query_edge_counts.groupby(["c_s", "r"]).head(3).sum() / num_edges
         hits_at_10_limit = community_query_edge_counts.groupby(["c_s", "r"]).head(10).sum() / num_edges
         community_query_counts = community_query_edge_counts.groupby(["c_s", "r"]).count()
-        community_query_edge_counts = community_query_edge_counts.to_frame().assign(rank=0, rrank=0)
-        for (c_s, r), c_t_count in community_query_counts.iteritems():
+        community_query_edge_counts = community_query_edge_counts.to_frame(name="c_t").assign(rank=0, rrank=0.0)
+        for (c_s, r), c_t_count in community_query_counts.items():
             community_query_edge_counts.loc[(c_s, r), "rank"] = np.arange(1, c_t_count + 1)
             community_query_edge_counts.loc[(c_s, r), "rrank"] = 1 / np.arange(1, c_t_count + 1)
         mr_limit = (community_query_edge_counts["c_t"] * community_query_edge_counts["rank"]).sum() / num_edges

src/research/MultiProxAn/src/analysis/rdkit_functions.py

@@ -1,7 +1,6 @@
 import numpy as np
 import torch
 import re
-import wandb
 try:
     from rdkit import Chem
     print("Found rdkit, all good")
@@ -316,19 +315,10 @@ def compute_molecular_metrics(molecule_list, train_smiles, dataset_info):
         fraction_mol_stable = molecule_stable / float(n_molecules)
         fraction_atm_stable = nr_stable_bonds / float(n_atoms)
         validity_dict = {'mol_stable': fraction_mol_stable, 'atm_stable': fraction_atm_stable}
-        if wandb.run:
-            wandb.log(validity_dict)
     else:
         validity_dict = {'mol_stable': -1, 'atm_stable': -1}
 
     metrics = BasicMolecularMetrics(dataset_info, train_smiles)
     rdkit_metrics = metrics.evaluate(molecule_list)
     all_smiles = rdkit_metrics[-1]
-    if wandb.run:
-        nc = rdkit_metrics[-2]
-        dic = {'Validity': rdkit_metrics[0][0], 'Relaxed Validity': rdkit_metrics[0][1],
-               'Uniqueness': rdkit_metrics[0][2], 'Novelty': rdkit_metrics[0][3],
-               'nc_max': nc['nc_max'], 'nc_mu': nc['nc_mu']}
-        wandb.log(dic)
-
     return validity_dict, rdkit_metrics, all_smiles

src/research/MultiProxAn/src/analysis/spectre_utils.py

@@ -3,7 +3,6 @@
 # Adapted from https://github.com/lrjconan/GRAN/ which in turn is adapted from https://github.com/JiaxuanYou/graph-generation
 #
 ###############################################################################
-import graph_tool.all as gt
 ##Navigate to the ./util/orca directory and compile orca.cpp
 # g++ -O2 -std=c++11 -o orca orca.cpp
 import os
@@ -12,18 +11,26 @@ import torch
 import torch.nn as nn
 import numpy as np
 import networkx as nx
-import subprocess as sp
-import concurrent.futures
-
-import pygsp as pg
-import secrets
-from string import ascii_uppercase, digits
-from datetime import datetime
-from scipy.linalg import eigvalsh
-from scipy.stats import chi2
-from src.analysis.dist_helper import compute_mmd, gaussian_emd, gaussian, emd, gaussian_tv, disc
-from torch_geometric.utils import to_networkx
-import wandb
+
+# Heavy metric-computation deps — optional, not needed for inference.
+# Deferred so checkpoint unpickling works without graph-tool / pyemd / pygsp.
+try:
+    import graph_tool.all as gt
+except ImportError:
+    gt = None
+try:
+    import subprocess as sp
+    import concurrent.futures
+    import pygsp as pg
+    import secrets
+    from string import ascii_uppercase, digits
+    from datetime import datetime
+    from scipy.linalg import eigvalsh
+    from scipy.stats import chi2
+    from src.analysis.dist_helper import compute_mmd, gaussian_emd, gaussian, emd, gaussian_tv, disc
+    from torch_geometric.utils import to_networkx
+except ImportError:
+    pass  # Metrics unavailable — inference still works
 
 PRINT_TIME = False
 __all__ = ['degree_stats', 'clustering_stats', 'orbit_stats_all', 'spectral_stats', 'eval_acc_lobster_graph']
@@ -778,8 +785,6 @@ class SpectreSamplingMetrics(nn.Module):
             degree = degree_stats(reference_graphs, networkx_graphs, is_parallel=False,
                                   compute_emd=self.compute_emd)
             to_log['degree'] = degree
-            if wandb.run:
-                wandb.run.summary['degree'] = degree
 
         # val_eigvals = [graph["eigval"][1:self.k + 1].cpu().detach().numpy() for graph in self.val]
         # train_eigvals = [graph["eigval"][1:self.k + 1].cpu().detach().numpy() for graph in self.train]
@@ -795,8 +800,6 @@ class SpectreSamplingMetrics(nn.Module):
                                      compute_emd=self.compute_emd)
 
             to_log['spectre'] = spectre
-            if wandb.run:
-                wandb.run.summary['spectre'] = spectre
 
         if 'clustering' in self.metrics_list:
             if local_rank == 0:
@@ -804,8 +807,6 @@ class SpectreSamplingMetrics(nn.Module):
             clustering = clustering_stats(reference_graphs, networkx_graphs, bins=100, is_parallel=False,
                                           compute_emd=self.compute_emd)
             to_log['clustering'] = clustering
-            if wandb.run:
-                wandb.run.summary['clustering'] = clustering
 
         if 'motif' in self.metrics_list:
             if local_rank == 0:
@@ -813,32 +814,24 @@ class SpectreSamplingMetrics(nn.Module):
             motif = motif_stats(reference_graphs, networkx_graphs, motif_type='4cycle', ground_truth_match=None, bins=100,
                                 compute_emd=self.compute_emd)
             to_log['motif'] = motif
-            if wandb.run:
-                wandb.run.summary['motif'] = motif
 
         if 'orbit' in self.metrics_list:
             if local_rank == 0:
                 print("Computing orbit stats...")
             orbit = orbit_stats_all(reference_graphs, networkx_graphs, compute_emd=self.compute_emd)
             to_log['orbit'] = orbit
-            if wandb.run:
-                wandb.run.summary['orbit'] = orbit
 
         if 'sbm' in self.metrics_list:
             if local_rank == 0:
                 print("Computing accuracy...")
             acc = eval_acc_sbm_graph(networkx_graphs, refinement_steps=100, strict=True)
             to_log['sbm_acc'] = acc
-            if wandb.run:
-                wandb.run.summary['sbmacc'] = acc
 
         if 'planar' in self.metrics_list:
             if local_rank ==0:
                 print('Computing planar accuracy...')
             planar_acc = eval_acc_planar_graph(networkx_graphs)
             to_log['planar_acc'] = planar_acc
-            if wandb.run:
-                wandb.run.summary['planar_acc'] = planar_acc
 
         if 'sbm' or 'planar' in self.metrics_list:
             if local_rank == 0:
@@ -853,8 +846,6 @@ class SpectreSamplingMetrics(nn.Module):
 
         if local_rank == 0:
             print("Sampling statistics", to_log)
-        if wandb.run:
-            wandb.log(to_log, commit=False)
 
     def reset(self):
         pass

src/research/MultiProxAn/src/analysis/visualization.py

@@ -8,7 +8,6 @@ import imageio
 import networkx as nx
 import numpy as np
 import rdkit.Chem
-import wandb
 import matplotlib.pyplot as plt
 
 
@@ -78,9 +77,6 @@ class MolecularVisualization:
             mol = self.mol_from_graphs(molecules[i][0].numpy(), molecules[i][1].numpy())
             try:
                 Draw.MolToFile(mol, file_path)
-                if wandb.run and log is not None:
-                    print(f"Saving {file_path} to wandb")
-                    wandb.log({log: wandb.Image(file_path)}, commit=True)
             except rdkit.Chem.KekulizeException:
                 print("Can't kekulize molecule")
 
@@ -115,10 +111,6 @@ class MolecularVisualization:
         imgs.extend([imgs[-1]] * 10)
         imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)
 
-        if wandb.run:
-            print(f"Saving {gif_path} to wandb")
-            wandb.log({"chain": wandb.Video(gif_path, fps=5, format="gif")}, commit=True)
-
         # draw grid image
         try:
             img = Draw.MolsToGridImage(mols, molsPerRow=20, subImgSize=(200, 200))
@@ -185,8 +177,6 @@ class NonMolecularVisualization:
             graph = self.to_networkx(graphs[i][0].numpy(), graphs[i][1].numpy())
             self.visualize_non_molecule(graph=graph, pos=None, path=file_path)
             im = plt.imread(file_path)
-            if wandb.run and log is not None:
-                wandb.log({log: [wandb.Image(im, caption=file_path)]})
 
     def visualize_chain(self, path, nodes_list, adjacency_matrix):
         # convert graphs to networkx
@@ -219,5 +209,3 @@ class NonMolecularVisualization:
         gif_path = os.path.join(os.path.dirname(path), '{}.gif'.format(path.split('/')[-1]))
         imgs.extend([imgs[-1]] * 10)
         imageio.mimsave(gif_path, imgs, subrectangles=True, duration=20)
-        if wandb.run:
-            wandb.log({'chain': [wandb.Video(gif_path, caption=gif_path, format="gif")]})

src/research/MultiProxAn/src/diffusion_model.py

@@ -6,7 +6,6 @@ import numpy as np
 import pytorch_lightning as pl
 import torch
 import torch.nn as nn
-import wandb
 from tqdm.auto import tqdm
 
 from diffusion.noise_schedule import PredefinedNoiseSchedule
@@ -141,8 +140,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
 
     def on_fit_start(self) -> None:
         self.train_iterations = len(self.trainer.datamodule.train_dataloader())
-        if self.local_rank == 0:
-            utils.setup_wandb(self.cfg)
 
     def on_train_epoch_start(self) -> None:
         self.start_epoch_time = time.time()
@@ -186,15 +183,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
         metrics = [self.val_nll.compute(), self.val_X_mse.compute(), self.val_E_mse.compute(),
                    self.val_y_mse.compute(), self.val_X_logp.compute(), self.val_E_logp.compute(),
                    self.val_y_logp.compute()]
-        if wandb.run:
-            wandb.log({"val/epoch_NLL": metrics[0],
-                       "val/X_mse": metrics[1],
-                       "val/E_mse": metrics[2],
-                       "val/y_mse": metrics[3],
-                       "val/X_logp": metrics[4],
-                       "val/E_logp": metrics[5],
-                       "val/y_logp": metrics[6]}, commit=False)
-
         print(f"Epoch {self.current_epoch}: Val NLL {metrics[0] :.2f} -- Val Atom type MSE {metrics[1] :.2f} -- ",
               f"Val Edge type MSE: {metrics[2] :.2f} -- Val Global feat. MSE {metrics[3] :.2f}",
               f"-- Val X Reconstruction loss {metrics[4] :.2f} -- Val E Reconstruction loss {metrics[5] :.2f}",
@@ -203,8 +191,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
         # Log val nll with default Lightning logger, so it can be monitored by checkpoint callback
         val_nll = metrics[0]
         self.log("val/epoch_NLL", val_nll, sync_dist=True)
-        if wandb.run:
-            wandb.log(self.log_info(), commit=False)
 
         if val_nll < self.best_val_nll:
             self.best_val_nll = val_nll
@@ -249,8 +235,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
         self.test_X_logp.reset()
         self.test_E_logp.reset()
         self.test_y_logp.reset()
-        if self.local_rank == 0:
-            utils.setup_wandb(self.cfg)
 
     def test_step(self, data, i):
         dense_data, node_mask = utils.to_dense(x=data.x, edge_index=data.edge_index, edge_attr=data.edge_attr,
@@ -277,19 +261,12 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
                     "test/X_logp": metrics[4],
                     "test/E_logp": metrics[5],
                     "test/y_logp": metrics[6]}
-        if wandb.run:
-            wandb.log(log_dict, commit=False)
-
         print(f"Epoch {self.current_epoch}: Test NLL {metrics[0] :.2f} -- Test Atom type MSE {metrics[1] :.2f} -- ",
               f"Test Edge type MSE: {metrics[2] :.2f} -- Test Global feat. MSE {metrics[3] :.2f}",
               f"-- Test X Reconstruction loss {metrics[4] :.2f} -- Test E Reconstruction loss {metrics[5] :.2f}",
               f"-- Test y Reconstruction loss {metrics[6] : .2f}\n")
 
         test_nll = metrics[0]
-        if wandb.run:
-            wandb.log({"test/epoch_NLL": test_nll}, commit=False)
-            wandb.log(self.log_info(), commit=False)
-
         print(f'Test loss: {test_nll :.4f}')
 
         samples_left_to_generate = self.cfg.general.final_model_samples_to_generate
@@ -320,9 +297,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
             samples_left_to_save -= to_save
             samples_left_to_generate -= to_generate
             chains_left_to_save -= chains_save
-        if wandb.run:
-            wandb.log({"test/time": total_eval_time}, commit=False)
-            wandb.run.summary['test_time'] = total_eval_time
         print(f'Test time: {total_eval_time :.4f} seconds')
 
         self.sampling_metrics.reset()
@@ -587,12 +561,6 @@ class LiftedDenoisingDiffusion(pl.LightningModule):
 
         nll = self.test_nll(nlls) if test else self.val_nll(nlls)  # Average over the batch
 
-        wandb.log({"kl prior": kl_prior.mean(),
-                   "Estimator loss terms": loss_all_t.mean(),
-                   "Loss term 0": loss_term_0,
-                   "log_pn": log_pN.mean(),
-                   'test_nll' if test else 'val_nll': nll},
-                  commit=False)
         return nll
 
     def forward(self, noisy_data, extra_data, node_mask):

src/research/MultiProxAn/src/diffusion_model_discrete.py

@@ -6,7 +6,6 @@ import pytorch_lightning as pl
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import wandb
 from tqdm import tqdm
 
 from diffusion.noise_schedule import DiscreteUniformTransition, MarginalUniformTransition, \
@@ -148,8 +147,6 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
     def on_fit_start(self) -> None:
         self.train_iterations = len(self.trainer.datamodule.train_dataloader())
         self.print("Size of the input features", self.Xdim, self.Edim, self.ydim)
-        if self.local_rank == 0:
-            utils.setup_wandb(self.cfg)
 
     def on_train_epoch_start(self) -> None:
         self.print("Starting train epoch...")
@@ -187,13 +184,6 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
     def on_validation_epoch_end(self) -> None:
         metrics = [self.val_nll.compute(), self.val_X_kl.compute() * self.T, self.val_E_kl.compute() * self.T,
                    self.val_X_logp.compute(), self.val_E_logp.compute()]
-        if wandb.run:
-            wandb.log({"val/epoch_NLL": metrics[0],
-                       "val/X_kl": metrics[1],
-                       "val/E_kl": metrics[2],
-                       "val/X_logp": metrics[3],
-                       "val/E_logp": metrics[4]}, commit=False)
-
         self.print(f"Epoch {self.current_epoch}: Val NLL {metrics[0] :.2f} -- Val Atom type KL {metrics[1] :.2f} -- ",
                    f"Val Edge type KL: {metrics[2] :.2f}")
 
@@ -242,8 +232,6 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
         self.test_E_kl.reset()
         self.test_X_logp.reset()
         self.test_E_logp.reset()
-        if self.local_rank == 0:
-            utils.setup_wandb(self.cfg)
 
     def test_step(self, data, i):
         dense_data, node_mask = utils.to_dense(data.x, data.edge_index, data.edge_attr, data.batch)
@@ -258,20 +246,10 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
         """ Measure likelihood on a test set and compute stability metrics. """
         metrics = [self.test_nll.compute(), self.test_X_kl.compute(), self.test_E_kl.compute(),
                    self.test_X_logp.compute(), self.test_E_logp.compute()]
-        if wandb.run:
-            wandb.log({"test/epoch_NLL": metrics[0],
-                       "test/X_kl": metrics[1],
-                       "test/E_kl": metrics[2],
-                       "test/X_logp": metrics[3],
-                       "test/E_logp": metrics[4]}, commit=False)
-
         self.print(f"Epoch {self.current_epoch}: Test NLL {metrics[0] :.2f} -- Test Atom type KL {metrics[1] :.2f} -- ",
                    f"Test Edge type KL: {metrics[2] :.2f}")
 
         test_nll = metrics[0]
-        if wandb.run:
-            wandb.log({"test/epoch_NLL": test_nll}, commit=False)
-
         self.print(f'Test loss: {test_nll :.4f}')
 
         samples_left_to_generate = self.cfg.general.final_model_samples_to_generate
@@ -304,9 +282,6 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
             samples_left_to_save -= to_save
             samples_left_to_generate -= to_generate
             chains_left_to_save -= chains_save
-        if wandb.run:
-            wandb.log({"test/time": total_eval_time}, commit=False)
-            wandb.run.summary['test_time'] = total_eval_time
         print(f'Test time: {total_eval_time :.4f} seconds')
 
         self.print("Saving the generated graphs")
@@ -531,12 +506,6 @@ class DiscreteDenoisingDiffusion(pl.LightningModule):
         # Update NLL metric object and return batch nll
         nll = (self.test_nll if test else self.val_nll)(nlls)  # Average over the batch
 
-        if wandb.run:
-            wandb.log({"kl prior": kl_prior.mean(),
-                       "Estimator loss terms": loss_all_t.mean(),
-                       "log_pn": log_pN.mean(),
-                       "loss_term_0": loss_term_0,
-                       'batch_test_nll' if test else 'val_nll': nll}, commit=False)
         return nll
 
     def forward(self, noisy_data, extra_data, node_mask):

src/research/MultiProxAn/src/metrics/molecular_metrics.py

@@ -6,7 +6,6 @@ from src.analysis.rdkit_functions import compute_molecular_metrics
 import torch
 from torchmetrics import Metric, MetricCollection
 from torch import Tensor
-import wandb
 import torch.nn as nn
 
 
@@ -25,8 +24,6 @@ class TrainMolecularMetrics(nn.Module):
                 to_log['train/' + key] = val.item()
             for key, val in self.train_bond_metrics.compute().items():
                 to_log['train/' + key] = val.item()
-            if wandb.run:
-                wandb.log(to_log, commit=False)
 
     def reset(self):
         for metric in [self.train_atom_metrics, self.train_bond_metrics]:
@@ -42,9 +39,6 @@ class TrainMolecularMetrics(nn.Module):
         for key, val in epoch_bond_metrics.items():
             to_log['train_epoch/epoch' + key] = val.item()
 
-        if wandb.run:
-            wandb.log(to_log, commit=False)
-
         for key, val in epoch_atom_metrics.items():
             epoch_atom_metrics[key] = f"{val.item() :.3f}"
         for key, val in epoch_bond_metrics.items():
@@ -135,18 +129,6 @@ class SamplingMolecularMetrics(nn.Module):
         edge_mae = self.edge_dist_mae.compute()
         valency_mae = self.valency_dist_mae.compute()
 
-        if wandb.run:
-            wandb.log(to_log, commit=False)
-            wandb.run.summary['Gen n distribution'] = generated_n_dist
-            wandb.run.summary['Gen node distribution'] = generated_node_dist
-            wandb.run.summary['Gen edge distribution'] = generated_edge_dist
-            wandb.run.summary['Gen valency distribution'] = generated_valency_dist
-
-            wandb.log({'basic_metrics/n_mae': n_mae,
-                       'basic_metrics/node_mae': node_mae,
-                       'basic_metrics/edge_mae': edge_mae,
-                       'basic_metrics/valency_mae': valency_mae}, commit=False)
-
         if local_rank == 0:
             print("Custom metrics computed.")
         if local_rank == 0:

src/research/MultiProxAn/src/metrics/molecular_metrics_discrete.py

@@ -1,7 +1,6 @@
 import torch
 from torchmetrics import Metric, MetricCollection
 from torch import Tensor
-import wandb
 import torch.nn as nn
 
 
@@ -167,8 +166,6 @@ class TrainMolecularMetricsDiscrete(nn.Module):
                 to_log['train/' + key] = val.item()
             for key, val in self.train_bond_metrics.compute().items():
                 to_log['train/' + key] = val.item()
-            if wandb.run:
-                wandb.log(to_log, commit=False)
 
     def reset(self):
         for metric in [self.train_atom_metrics, self.train_bond_metrics]:
@@ -183,9 +180,6 @@ class TrainMolecularMetricsDiscrete(nn.Module):
             to_log['train_epoch/' + key] = val.item()
         for key, val in epoch_bond_metrics.items():
             to_log['train_epoch/' + key] = val.item()
-        if wandb.run:
-            wandb.log(to_log, commit=False)
-
         for key, val in epoch_atom_metrics.items():
             epoch_atom_metrics[key] = val.item()
         for key, val in epoch_bond_metrics.items():

src/research/MultiProxAn/src/metrics/train_metrics.py

@@ -3,7 +3,6 @@ from torch import Tensor
 import torch.nn as nn
 from torchmetrics import Metric, MeanSquaredError, MetricCollection
 import time
-import wandb
 from src.metrics.abstract_metrics import SumExceptBatchMetric, SumExceptBatchMSE, SumExceptBatchKL, CrossEntropyMetric, \
     ProbabilityMetric, NLL
 
@@ -36,8 +35,6 @@ class TrainLoss(nn.Module):
                       'train_loss/node_MSE': self.train_node_mse.compute(),
                       'train_loss/edge_MSE': self.train_edge_mse.compute(),
                       'train_loss/y_mse': self.train_y_mse.compute()}
-            if wandb.run:
-                wandb.log(to_log, commit=True)
 
         return mse
 
@@ -53,8 +50,6 @@ class TrainLoss(nn.Module):
         to_log = {"train_epoch/epoch_X_mse": epoch_node_mse,
                   "train_epoch/epoch_E_mse": epoch_edge_mse,
                   "train_epoch/epoch_y_mse": epoch_y_mse}
-        if wandb.run:
-            wandb.log(to_log)
         return to_log
 
 
@@ -101,8 +96,6 @@ class TrainLossDiscrete(nn.Module):
                       "train_loss/X_CE": self.node_loss.compute() if true_X.numel() > 0 else -1,
                       "train_loss/E_CE": self.edge_loss.compute() if true_E.numel() > 0 else -1,
                       "train_loss/y_CE": self.y_loss.compute() if true_y.numel() > 0 else -1}
-            if wandb.run:
-                wandb.log(to_log, commit=True)
         return loss_X + self.lambda_train[0] * loss_E + self.lambda_train[1] * loss_y
 
     def reset(self):
@@ -117,9 +110,6 @@ class TrainLossDiscrete(nn.Module):
         to_log = {"train_epoch/x_CE": epoch_node_loss,
                   "train_epoch/E_CE": epoch_edge_loss,
                   "train_epoch/y_CE": epoch_y_loss}
-        if wandb.run:
-            wandb.log(to_log, commit=False)
-
         return to_log
 
 

src/research/MultiProxAn/src/utils.py

@@ -4,7 +4,6 @@ from omegaconf import OmegaConf, open_dict
 from torch_geometric.utils import to_dense_adj, to_dense_batch
 import torch
 import omegaconf
-import wandb
 
 
 def create_folders(args):
@@ -131,9 +130,3 @@ class PlaceHolder:
         return self
 
 
-def setup_wandb(cfg):
-    config_dict = omegaconf.OmegaConf.to_container(cfg, resolve=True, throw_on_missing=True)
-    kwargs = {'name': cfg.general.name, 'project': f'graph_ddm_{cfg.dataset.name}', 'config': config_dict,
-              'settings': wandb.Settings(_disable_stats=True), 'reinit': True, 'mode': cfg.general.wandb}
-    wandb.init(**kwargs)
-    wandb.save('*.txt')