diff --git a/pdebench/data_gen/data_gen_NLE/config/config.yaml b/pdebench/data_gen/data_gen_NLE/config/config.yaml
index 25f1620..89a88d5 100644
--- a/pdebench/data_gen/data_gen_NLE/config/config.yaml
+++ b/pdebench/data_gen/data_gen_NLE/config/config.yaml
@@ -9,7 +9,7 @@ hydra:
     dir: .
 
 args:
-  type: "ReacDiff"  # "advection"/"ReacDiff"/"burgers"/"CFD"
+  type: "ReacDiff" # "advection"/"ReacDiff"/"burgers"/"CFD"
   dim: 1
   bd: "periodic"
   nbatch: 1000
diff --git a/pdebench/models/fno/train.py b/pdebench/models/fno/train.py
index 2381c45..d299d1f 100644
--- a/pdebench/models/fno/train.py
+++ b/pdebench/models/fno/train.py
@@ -6,11 +6,11 @@
 
 import numpy as np
 import torch
-import torch.nn as nn
 
 from pdebench.models.fno.fno import FNO1d, FNO2d, FNO3d
 from pdebench.models.fno.utils import FNODatasetMult, FNODatasetSingle
 from pdebench.models.metrics import metrics
+from torch import nn
 
 # torch.manual_seed(0)
 # np.random.seed(0)