lint

MARBLE/main.py

@@ -4,16 +4,19 @@
 import warnings
 from datetime import datetime
 from pathlib import Path
-import yaml
-from tqdm import tqdm
 
 import torch
 import torch.nn.functional as F
 import torch.optim as opt
+import yaml
 from torch import nn
 from torch_geometric.nn import MLP
+from tqdm import tqdm
 
-from MARBLE import dataloader, geometry, layers, utils
+from MARBLE import dataloader
+from MARBLE import geometry
+from MARBLE import layers
+from MARBLE import utils
 
 
 class net(nn.Module):

MARBLE/plotting.py

@@ -1,12 +1,13 @@
 """Plotting module."""
 import os
 from pathlib import Path
-import torch
+
 import matplotlib
 import matplotlib.pyplot as plt
 import networkx as nx
 import numpy as np
 import seaborn as sns
+import torch
 from matplotlib import gridspec
 from matplotlib.colors import LinearSegmentedColormap
 from matplotlib.patches import FancyArrowPatch

MARBLE/preprocessing.py

@@ -125,7 +125,7 @@ def _compute_geometric_objects(
         gauges (nxdxd): local gauges at all points
         par (dict): updated dictionary of parameters
         local_gauges: whether to use local gauges
-        
+
     """
     n, dim_emb = data.pos.shape
     dim_signal = data.x.shape[1]
@@ -184,4 +184,4 @@ def _compute_geometric_objects(
     ]
     data.L, data.Lc, data.gauges, data.local_gauges = L, Lc, gauges, local_gauges
 
-    return data
+    return data