diff --git a/examples/problems/elliptic_optimal_control.py b/examples/problems/elliptic_optimal_control.py
index f110626..db8d056 100644
--- a/examples/problems/elliptic_optimal_control.py
+++ b/examples/problems/elliptic_optimal_control.py
@@ -1,4 +1,3 @@
-import numpy as np
 import torch
 from pina.problem import Problem
 from pina.segment import Segment
@@ -39,11 +38,8 @@ class EllipticOptimalControl(Problem2D):
             'gamma3': {'location': Segment((xmax, ymax), (xmin, ymax)), 'func': nil_dirichlet},
             'gamma4': {'location': Segment((xmin, ymax), (xmin, ymin)), 'func': nil_dirichlet},
             'D1': {'location': Cube([[xmin, xmax], [ymin, ymax]]), 'func': [term1, term2, term3]},
-            #'D2': {'location': Cube([[0, 1], [0, 1]]), 'func': term2},
-            #'D3': {'location': Cube([[0, 1], [0, 1]]), 'func': term3}
         }
 
         self.input_variables = ['x1', 'x2']
         self.output_variables = ['u', 'p', 'y']
         self.spatial_domain = Cube([[xmin, xmax], [xmin, xmax]])
-
diff --git a/pina/__init__.py b/pina/__init__.py
index c1ce52a..d764327 100644
--- a/pina/__init__.py
+++ b/pina/__init__.py
@@ -1,16 +1,16 @@
 __all__ = [
     'PINN',
-    'ParametricPINN',
     'LabelTensor',
     'Plotter',
     'Condition',
-    'Span'
+    'Span',
+    'Location',
 ]
 
 from .meta import *
 from .label_tensor import LabelTensor
 from .pinn import PINN
-#from .ppinn import ParametricPINN
 from .plotter import Plotter
 from .span import Span
 from .condition import Condition
+from .location import Location
diff --git a/pina/model/__init__.py b/pina/model/__init__.py
index ccdc8a8..5f55123 100644
--- a/pina/model/__init__.py
+++ b/pina/model/__init__.py
@@ -1,8 +1,8 @@
 __all__ = [
     'FeedForward',
-    'MultiFeedForward'
+    'MultiFeedForward',
     'DeepONet',
-    'Network'
+    'Network',
 ]
 
 from .feed_forward import FeedForward
diff --git a/pina/pinn.py b/pina/pinn.py
index e6cecbf..860c464 100644
--- a/pina/pinn.py
+++ b/pina/pinn.py
@@ -1,8 +1,9 @@
-from .problem import AbstractProblem
+""" Module for PINN """
 import torch
-import matplotlib.pyplot as plt
-import numpy as np
-from pina.label_tensor import LabelTensor
+
+from .problem import AbstractProblem
+from .label_tensor import LabelTensor
+
 torch.pi = torch.acos(torch.zeros(1)).item() * 2 # which is 3.1415927410125732