diff --git a/bindings/pyroot/pythonizations/CMakeLists.txt b/bindings/pyroot/pythonizations/CMakeLists.txt
index 9dc781dd31a15..b046c3435f8d7 100644
--- a/bindings/pyroot/pythonizations/CMakeLists.txt
+++ b/bindings/pyroot/pythonizations/CMakeLists.txt
@@ -68,7 +68,7 @@ if(tmva)
 endif()
 
 if(PYTHON_VERSION_STRING_Development_Main VERSION_GREATER_EQUAL 3.8 AND dataframe)
-    list(APPEND PYROOT_EXTRA_PY3_SOURCE
+list(APPEND PYROOT_EXTRA_PY3_SOURCE
      ROOT/_pythonization/_tmva/_batchgenerator.py)
 endif()
 
diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/__init__.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/__init__.py
index 2884acf0fb60d..2603fa040f224 100644
--- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/__init__.py
+++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/__init__.py
@@ -41,8 +41,8 @@ def inject_rbatchgenerator(ns):
             setattr(ns.Experimental, func_name, python_func)
 
         return ns
-    
-from ._gnn import RModel_GNN, RModel_GraphIndependent
+
+    from ._gnn import RModel_GNN, RModel_GraphIndependent
 
 hasRDF = gSystem.GetFromPipe("root-config --has-dataframe") == "yes"
 if hasRDF:
diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_gnn.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_gnn.py
index d4aa562126112..35ca123e18dbf 100644
--- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_gnn.py
+++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_tmva/_gnn.py
@@ -14,7 +14,7 @@
 import sys
 from cppyy import gbl as gbl_namespace
 
-if sys.version_info < (3, 7):
+if sys.version_info < (3, 8):
     raise RuntimeError("GNN Pythonizations are only supported in Python3")
 
 
@@ -29,7 +29,7 @@ def getActivationFunction(model):
 
     Returns:
         The activation function enum value.
-    """    
+    """
     function = model._activation.__name__
     if function == 'relu':
         return gbl_namespace.TMVA.Experimental.SOFIE.Activation.RELU
@@ -80,7 +80,7 @@ def add_layer_norm(gin, module_layer, function_target):
     else:
         model_block = gin.globals_update_block
     axis = module_layer._axis
-    eps  = module_layer._eps 
+    eps  = module_layer._eps
     stash_type = 1
     name_x = model_block.GetFunctionBlock().GetOutputTensorNames()[0]
     name_bias = module_layer.offset.name
@@ -90,7 +90,7 @@ def add_layer_norm(gin, module_layer, function_target):
     current_output_tensors = model_block.GetFunctionBlock().GetOutputTensorNames()
     new_output_tensors = gbl_namespace.std.vector['std::string']()
     new_output_tensors.push_back(name_Y)
-    model_block.GetFunctionBlock().AddOutputTensorNameList(new_output_tensors)    
+    model_block.GetFunctionBlock().AddOutputTensorNameList(new_output_tensors)
 
 def add_weights(gin, weights, function_target):
     """
@@ -133,12 +133,12 @@ def add_aggregate_function(gin, reducer, relation):
         agg = gbl_namespace.TMVA.Experimental.SOFIE.RFunction_Mean()
         gin.createAggregateFunction[gbl_namespace.TMVA.Experimental.SOFIE.RFunction_Mean](agg, relation)
     else:
-        raise RuntimeError("Invalid aggregate function for reduction")    
+        raise RuntimeError("Invalid aggregate function for reduction")
 
 
 def add_update_function(gin, component_model, graph_type, function_target):
     """
-    Add update function for respective function target, either of nodes, edges or globals 
+    Add update function for respective function target, either of nodes, edges or globals
     based on the supplied component_model
 
     Parameters:
@@ -164,7 +164,7 @@ def add_update_function(gin, component_model, graph_type, function_target):
 
 
 
-class RModel_GNN: 
+class RModel_GNN:
     """
     Wrapper class for graph_nets' GNN model;s parsing and inference generation
 
@@ -199,21 +199,21 @@ def ParseFromMemory(graph_module, graph_data, filename = "gnn_network"):
         gin.num_global_features = len(graph_data['globals'])
 
         gin.filename = filename
-        
+
         # adding the node update function
         node_model = graph_module._node_block._node_model
-        add_update_function(gin, node_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN, 
+        add_update_function(gin, node_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN,
                                          gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.NODES)
 
         # adding the edge update function
         edge_model = graph_module._edge_block._edge_model
-        add_update_function(gin, edge_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN, 
+        add_update_function(gin, edge_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN,
                                              gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.EDGES)
 
         # adding the global update function
         global_model = graph_module._global_block._global_model
-        add_update_function(gin, global_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN, 
-                                             gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.GLOBALS) 
+        add_update_function(gin, global_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GNN,
+                                             gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.GLOBALS)
 
         # adding edge-node aggregate function
         add_aggregate_function(gin, graph_module._node_block._received_edges_aggregator._reducer.__qualname__, gbl_namespace.TMVA.Experimental.SOFIE.FunctionRelation.NODES_EDGES)
@@ -274,18 +274,18 @@ def ParseFromMemory(graph_module, graph_data, filename = "graph_independent_netw
 
         # adding the node update function
         node_model = graph_module._node_model._model
-        add_update_function(gin, node_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent, 
+        add_update_function(gin, node_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent,
                                          gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.NODES)
 
         # adding the edge update function
         edge_model = graph_module._edge_model._model
-        add_update_function(gin, edge_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent, 
+        add_update_function(gin, edge_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent,
                                              gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.EDGES)
 
         # adding the global update function
         global_model = graph_module._global_model._model
-        add_update_function(gin, global_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent, 
-                                             gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.GLOBALS) 
+        add_update_function(gin, global_model, gbl_namespace.TMVA.Experimental.SOFIE.GraphType.GraphIndependent,
+                                             gbl_namespace.TMVA.Experimental.SOFIE.FunctionTarget.GLOBALS)
 
         graph_independent_model = gbl_namespace.TMVA.Experimental.SOFIE.RModel_GraphIndependent(gin)
         blas_routines = gbl_namespace.std.vector['std::string']()
diff --git a/bindings/pyroot/pythonizations/test/CMakeLists.txt b/bindings/pyroot/pythonizations/test/CMakeLists.txt
index b31e4f613c39e..c6e526e0507c5 100644
--- a/bindings/pyroot/pythonizations/test/CMakeLists.txt
+++ b/bindings/pyroot/pythonizations/test/CMakeLists.txt
@@ -128,12 +128,16 @@ endif()
 # SOFIE-GNN pythonizations
 if (dataframe AND tmva)
     if(NOT MSVC OR CMAKE_SIZEOF_VOID_P EQUAL 4 OR win_broken_tests AND PYTHON_VERSION_MAJOR_Development_Main EQUAL 3 )
-        ROOT_ADD_PYUNITTEST(pyroot_pyz_sofie_gnn sofie_gnn.py PYTHON_DEPS numpy sonnet graph_nets)
+        find_python_module(sonnet QUIET)
+        find_python_module(graph_nets QUIET)
+        if (PY_SONNET_FOUND AND PY_GRAPH_NETS_FOUND)
+            ROOT_ADD_PYUNITTEST(pyroot_pyz_sofie_gnn sofie_gnn.py PYTHON_DEPS numpy sonnet graph_nets)
+        endif()
     endif()
 endif()
 
 # RTensor pythonizations
-if (tmva AND sofie)
+if (tmva)
     if(NOT MSVC OR CMAKE_SIZEOF_VOID_P EQUAL 4 OR win_broken_tests)
         ROOT_ADD_PYUNITTEST(pyroot_pyz_rtensor rtensor.py PYTHON_DEPS numpy)
     endif()
diff --git a/bindings/pyroot/pythonizations/test/sofie_gnn.py b/bindings/pyroot/pythonizations/test/sofie_gnn.py
index aa861efa90d0a..a86bdfdf83194 100644
--- a/bindings/pyroot/pythonizations/test/sofie_gnn.py
+++ b/bindings/pyroot/pythonizations/test/sofie_gnn.py
@@ -4,12 +4,13 @@
 import numpy as np
 from numpy.testing import assert_almost_equal
 
-if np.__version__ > "1.19":
-    raise RuntimeError(f"This test requires NumPy version 1.19 or lower")
+if np.__version__ >= "1.20" and np.__version__ < "1.24":
+    raise RuntimeError(f"This test requires NumPy version <=1.19 or >=1.24")
 
 import graph_nets as gn
 from graph_nets import utils_tf
 import sonnet as snt
+import os
 
 
 
@@ -21,9 +22,22 @@ def get_graph_data_dict(num_nodes, num_edges, GLOBAL_FEATURE_SIZE=2, NODE_FEATUR
       "nodes": np.random.rand(num_nodes, NODE_FEATURE_SIZE).astype(np.float32),
       "edges": np.random.rand(num_edges, EDGE_FEATURE_SIZE).astype(np.float32),
       "senders": np.random.randint(num_nodes, size=num_edges, dtype=np.int32),
-      "receivers": np.random.randint(num_nodes, size=num_edges, dtype=np.int32),
+      "receivers": np.random.randint(num_nodes, size=num_edges, dtype=np.int32)
   }
 
+def resize_graph_data(input_data, GLOBAL_FEATURE_SIZE=2, NODE_FEATURE_SIZE=2, EDGE_FEATURE_SIZE=2) :
+    n = input_data["nodes"]
+    e = input_data["edges"]
+    s = input_data["senders"]
+    r = input_data["receivers"]
+    return {
+      "globals" : np.zeros((GLOBAL_FEATURE_SIZE )),
+      "nodes"   : np.zeros((n.shape[0],NODE_FEATURE_SIZE )),
+      "edges"   : np.zeros((e.shape[0],EDGE_FEATURE_SIZE )),
+      "senders" : s, "receivers" : r
+    }
+
+LATENT_SIZE = 2
 def make_mlp_model():
   """Instantiates a new MLP, followed by LayerNorm.
 
@@ -34,7 +48,7 @@ def make_mlp_model():
     A Sonnet module which contains the MLP and LayerNorm.
   """
   return snt.Sequential([
-      snt.nets.MLP([2,2], activate_final=True),
+      snt.nets.MLP([LATENT_SIZE]*2, activate_final=True),
       snt.LayerNorm(axis=-1, create_offset=True, create_scale=True)
   ])
 
@@ -48,9 +62,9 @@ class MLPGraphIndependent(snt.Module):
   def __init__(self, name="MLPGraphIndependent"):
     super(MLPGraphIndependent, self).__init__(name=name)
     self._network = gn.modules.GraphIndependent(
-        edge_model_fn = lambda: snt.nets.MLP([2,2], activate_final=True),
-        node_model_fn = lambda: snt.nets.MLP([2,2], activate_final=True),
-        global_model_fn = lambda: snt.nets.MLP([2,2], activate_final=True))
+        edge_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*2, activate_final=True),
+        node_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*2, activate_final=True),
+        global_model_fn = lambda: snt.nets.MLP([LATENT_SIZE]*2, activate_final=True))
 
   def __call__(self, inputs):
     return self._network(inputs)
@@ -69,12 +83,19 @@ def __init__(self, name="MLPGraphNetwork"):
   def __call__(self, inputs):
     return self._network(inputs)
 
+def PrintGData(data, printShape = True):
+  n = data.nodes.numpy()
+  e = data.edges.numpy()
+  g = data.globals.numpy()
+  if (printShape) :
+    print("GNet data ... shapes",n.shape,e.shape,g.shape)
+  print(" node data", n.reshape(n.size,))
+  print(" edge data", e.reshape(e.size,))
+  print(" global data",g.reshape(g.size,))
+
 class EncodeProcessDecode(snt.Module):
 
   def __init__(self,
-               edge_output_size=None,
-               node_output_size=None,
-               global_output_size=None,
                name="EncodeProcessDecode"):
     super(EncodeProcessDecode, self).__init__(name=name)
     self._encoder = MLPGraphIndependent()
@@ -103,12 +124,15 @@ def test_parse_gnn(self):
         Test that parsed GNN model from a graphnets model generates correct
         inference code
         '''
+
+        print('\nRun Graph parsing test')
+
         GraphModule = gn.modules.GraphNetwork(
             edge_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True),
             node_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True),
             global_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True))
 
-        GraphData = get_graph_data_dict(2,1)
+        GraphData = get_graph_data_dict(2,1,2,2,2)
         input_graphs = utils_tf.data_dicts_to_graphs_tuple([GraphData])
         output = GraphModule(input_graphs)
 
@@ -135,18 +159,26 @@ def test_parse_gnn(self):
         assert_almost_equal(output_edge_data, np.asarray(input_data.edge_data))
         assert_almost_equal(output_global_data, np.asarray(input_data.global_data))
 
+        fname = "gnn_network"
+        os.remove(fname + '.dat')
+        os.remove(fname + '.hxx')
+
 
     def test_parse_graph_independent(self):
         '''
         Test that parsed GraphIndependent model from a graphnets model generates correct
         inference code
         '''
+
+        print('\nRun Graph Independent parsing test')
+
+
         GraphModule = gn.modules.GraphIndependent(
             edge_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True),
             node_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True),
             global_model_fn=lambda: snt.nets.MLP([2,2], activate_final=True))
 
-        GraphData = get_graph_data_dict(2,1)
+        GraphData = get_graph_data_dict(2,1,2,2,2)
         input_graphs = utils_tf.data_dicts_to_graphs_tuple([GraphData])
         output = GraphModule(input_graphs)
 
@@ -173,27 +205,46 @@ def test_parse_graph_independent(self):
         assert_almost_equal(output_edge_data, np.asarray(input_data.edge_data))
         assert_almost_equal(output_global_data, np.asarray(input_data.global_data))
 
+        fname = "graph_independent_network"
+        os.remove(fname + '.dat')
+        os.remove(fname + '.hxx')
+
     def test_lhcb_toy_inference(self):
         '''
         Test that parsed stack of SOFIE GNN and GraphIndependent modules generate the correct
         inference code
         '''
+
+        print('Run LHCb test')
+
         # Instantiating EncodeProcessDecode Model
-        ep_model = EncodeProcessDecode(2,2,2)
+
+        #number of features for node. edge, globals
+        nsize = 3
+        esize = 3
+        gsize = 2
+        lsize = LATENT_SIZE  #hard-coded latent size in definition of GNET model (for node edge and globals)
+
+        ep_model = EncodeProcessDecode()
 
         # Initializing randomized input data
-        GraphData = get_graph_data_dict(2,1)
-        input_graphs = utils_tf.data_dicts_to_graphs_tuple([GraphData])
+        InputGraphData = get_graph_data_dict(2,1, gsize, nsize, esize)
+        input_graphs = utils_tf.data_dicts_to_graphs_tuple([InputGraphData])
+
+        # Make data for core networks (number of features for node/edge global is 2 * lsize)
+        CoreGraphData = resize_graph_data(InputGraphData, 2 * lsize, 2 * lsize, 2 * lsize)
+
+
+        OutputGraphData = resize_graph_data(InputGraphData, lsize, lsize, lsize)
 
-        # Initializing randomized input data for core
-        CoreGraphData = get_graph_data_dict(2, 1, 4, 4, 4)
-        input_graphs_2 = utils_tf.data_dicts_to_graphs_tuple([CoreGraphData])
 
         # Collecting output from GraphNets model stack
         output_gn = ep_model(input_graphs, 2)
 
+        print("senders and receivers ",InputGraphData['senders'],InputGraphData['receivers'])
+
         # Declaring sofie models
-        encoder = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._encoder._network, GraphData, filename = "encoder")
+        encoder = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._encoder._network, InputGraphData, filename = "encoder")
         encoder.Generate()
         encoder.OutputGenerated()
 
@@ -201,11 +252,11 @@ def test_lhcb_toy_inference(self):
         core.Generate()
         core.OutputGenerated()
 
-        decoder = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._decoder._network, GraphData, filename = "decoder")
+        decoder = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._decoder._network, OutputGraphData, filename = "decoder")
         decoder.Generate()
         decoder.OutputGenerated()
 
-        output_transform = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._output_transform._network, GraphData, filename = "output_transform")
+        output_transform = ROOT.TMVA.Experimental.SOFIE.RModel_GraphIndependent.ParseFromMemory(ep_model._output_transform._network, OutputGraphData, filename = "output_transform")
         output_transform.Generate()
         output_transform.OutputGenerated()
 
@@ -222,14 +273,18 @@ def test_lhcb_toy_inference(self):
 
         # Preparing the input data for running inference on sofie
         input_data = ROOT.TMVA.Experimental.SOFIE.GNN_Data()
-        input_data.node_data = ROOT.TMVA.Experimental.AsRTensor(GraphData['nodes'])
-        input_data.edge_data = ROOT.TMVA.Experimental.AsRTensor(GraphData['edges'])
-        input_data.global_data = ROOT.TMVA.Experimental.AsRTensor(GraphData['globals'])
+        input_data.node_data = ROOT.TMVA.Experimental.AsRTensor(InputGraphData['nodes'])
+        input_data.edge_data = ROOT.TMVA.Experimental.AsRTensor(InputGraphData['edges'])
+        input_data.global_data = ROOT.TMVA.Experimental.AsRTensor(InputGraphData['globals'])
+
+        output_gn = ep_model(input_graphs, 2)
 
         # running inference on sofie
-        encoder_session.infer(input_data)
-        latent0 = CopyData(input_data)
-        latent = input_data
+        data = CopyData(input_data)
+
+        encoder_session.infer(data)
+        latent0 = CopyData(data)
+        latent = data
         output_ops = []
         for _ in range(2):
             core_input = ROOT.TMVA.Experimental.SOFIE.Concatenate(latent0, latent, axis=1)
@@ -246,10 +301,16 @@ def test_lhcb_toy_inference(self):
           output_global_data = output_gn[i].globals.numpy().flatten()
 
           assert_almost_equal(output_node_data, np.asarray(output_ops[i].node_data))
+
           assert_almost_equal(output_edge_data, np.asarray(output_ops[i].edge_data))
-          assert_almost_equal(output_global_data, np.asarray(output_ops[i].global_data))
 
+          assert_almost_equal(output_global_data, np.asarray(output_ops[i].global_data))
 
+        #remove header files after being used
+        filesToRemove = ['core','encoder','decoder','output_transform']
+        for fname in filesToRemove:
+          os.remove(fname + '.hxx')
+          os.remove(fname + '.dat')
 
 
 if __name__ == '__main__':
diff --git a/tmva/pymva/test/EmitFromKeras.cxx b/tmva/pymva/test/EmitFromKeras.cxx
index 9c84d42fe1fc4..4d6291e97f1b7 100644
--- a/tmva/pymva/test/EmitFromKeras.cxx
+++ b/tmva/pymva/test/EmitFromKeras.cxx
@@ -4,7 +4,6 @@
 //           The program is run when the target 'TestRModelParserKeras' is built.
 //           The program generates the required .hxx file after parsing a Keras .keras file into a RModel object.
 
-#include "TMVA/RModel_Base.hxx"
 #include "TMVA/RModel.hxx"
 #include "TMVA/RModelParser_Keras.h"
 
diff --git a/tmva/sofie/CMakeLists.txt b/tmva/sofie/CMakeLists.txt
index d44e238e30bca..abe4dd5540a97 100644
--- a/tmva/sofie/CMakeLists.txt
+++ b/tmva/sofie/CMakeLists.txt
@@ -58,6 +58,10 @@ ROOT_STANDARD_LIBRARY_PACKAGE(ROOTTMVASofie
     src/RModel.cxx
     src/RModel_GNN.cxx
     src/RModel_GraphIndependent.cxx
+    src/RFunction.cxx
+    src/RFunction_MLP.cxx
+    src/RFunction_Mean.cxx
+    src/RFunction_Sum.cxx
     src/SOFIE_common.cxx
   DEPENDENCIES
     TMVA
diff --git a/tmva/sofie/inc/TMVA/FunctionList.hxx b/tmva/sofie/inc/TMVA/FunctionList.hxx
index fa0e8f7bf72b9..eebadd3aaba70 100644
--- a/tmva/sofie/inc/TMVA/FunctionList.hxx
+++ b/tmva/sofie/inc/TMVA/FunctionList.hxx
@@ -3,4 +3,4 @@
 
 // Aggregate functions
 #include "TMVA/RFunction_Sum.hxx"
-#include "TMVA/RFunction_Mean.hxx"
\ No newline at end of file
+#include "TMVA/RFunction_Mean.hxx"
diff --git a/tmva/sofie/inc/TMVA/RFunction.hxx b/tmva/sofie/inc/TMVA/RFunction.hxx
index 2aff022b890e3..9d4f75d45f701 100644
--- a/tmva/sofie/inc/TMVA/RFunction.hxx
+++ b/tmva/sofie/inc/TMVA/RFunction.hxx
@@ -1,145 +1,81 @@
 #ifndef TMVA_SOFIE_RFUNCTION
 #define TMVA_SOFIE_RFUNCTION
 
-#include <any>
-#include "TMVA/RModel_GNN.hxx"
-#include <iostream>
+#include "TMVA/RModel_Base.hxx"
+#include "TMVA/SOFIE_common.hxx"
 
-namespace TMVA{
-namespace Experimental{
-namespace SOFIE{
+#include <memory>
+#include <string>
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
 
 class RModel;
 
 
-class RFunction{
-    protected:
-        std::string fFuncName;
-        FunctionType fType;
-    public:
-        RFunction(){}
-        virtual ~RFunction(){}
-        FunctionType GetFunctionType(){
-                return fType;
-        }
+class RFunction {
+protected:
+    std::string fFuncName;
+    FunctionType fType;
+public:
+    RFunction() {}
+    virtual ~RFunction() {}
+    FunctionType GetFunctionType() {
+        return fType;
+    }
 
-        RFunction(std::string funcName, FunctionType type):
-                fFuncName(UTILITY::Clean_name(funcName)),fType(type){}
+    RFunction(std::string funcName, FunctionType type):
+        fFuncName(UTILITY::Clean_name(funcName)),fType(type) {}
 
 };
 
-class RFunction_Update: public RFunction{
-        protected:
-                std::shared_ptr<RModel> function_block;
-                FunctionTarget fTarget;
-                GraphType fGraphType;
-                std::vector<std::string> fInputTensors;
-                std::vector<ROperator*> fAddlOp;  // temporary vector to store pointer that will be moved in a unique_ptr
-
-        public:
-        virtual ~RFunction_Update(){}
-        RFunction_Update(){}
-                RFunction_Update(FunctionTarget target, GraphType gType): fTarget(target), fGraphType(gType){
-                        switch(target){
-                                case FunctionTarget::EDGES:{
-                                        fFuncName = "edge_update";
-                                        break;
-                                }
-                                case FunctionTarget::NODES: {
-                                        fFuncName = "node_update";
-                                        break;
-                                }
-                                case FunctionTarget::GLOBALS: {
-                                        fFuncName = "global_update";
-                                        break;
-                                }
-                                default:
-                                        throw std::runtime_error("Invalid target for Update function");
-                        }
-                        fType = FunctionType::UPDATE;
-                        function_block = std::make_unique<RModel>(fFuncName);
-
-                        if(fGraphType == GraphType::GNN){
-                                if(fTarget == FunctionTarget::EDGES){
-                                        fInputTensors = {"edge","receiver","sender","global"};
-                                } else if(fTarget == FunctionTarget::NODES || fTarget == FunctionTarget::GLOBALS){
-                                        fInputTensors = {"edge","node","global"};
-                                }
-
-                        } else if(fGraphType == GraphType::GraphIndependent){
-                                if(fTarget == FunctionTarget::EDGES){
-                                        fInputTensors = {"edge"};
-                                } else if(fTarget == FunctionTarget::NODES){
-                                        fInputTensors = {"node"};
-                                } else {
-                                        fInputTensors = {"global"};
-                                }
-                        }
-                }
-
-                virtual void AddInitializedTensors(const std::vector<std::vector<std::string>>&){};
-                virtual void Initialize(){};
-                virtual void AddLayerNormalization(int, float, size_t, const std::string&,
-                                    const std::string&, const std::string&, const std::string&){};
-                void AddInputTensors(const std::vector<std::vector<std::size_t>>& fInputShape){
-                        for(long unsigned int i=0; i<fInputShape.size(); ++i){
-                                function_block->AddInputTensorInfo(fInputTensors[i],ETensorType::FLOAT, fInputShape[i]);
-                                function_block->AddInputTensorName(fInputTensors[i]);
-                        }
-                }
-                std::shared_ptr<RModel> GetFunctionBlock(){
-                        return function_block;
-                }
-
-                std::string GenerateModel(const std::string& filename, long read_pos=0, long block_size=1){
-                        function_block->SetFilename(filename);
-                        // use batch size as b;lock size in RModel::generate
-                        function_block->Generate(Options::kGNNComponent,block_size,read_pos);
-                        std::string modelGenerationString;
-                        modelGenerationString = "\n//--------- GNN_Update_Function---"+fFuncName+"\n"+function_block->ReturnGenerated();
-                        return modelGenerationString;
-                }
-                std::string Generate(const std::vector<std::string>& inputPtrs){
-                        std::string inferFunc = fFuncName+".infer(";
-                        for(auto&it : inputPtrs){
-                                inferFunc+=it;
-                                inferFunc+=",";
-                        }
-                        inferFunc.pop_back();
-                        inferFunc+=");";
-                        return inferFunc;
-                }
-                FunctionTarget GetFunctionTarget(){
-                        return fTarget;
-                }
+class RFunction_Update: public RFunction {
+protected:
+    std::shared_ptr<RModel> function_block;
+    FunctionTarget fTarget;
+    GraphType fGraphType;
+    std::vector<std::string> fInputTensors;
+    std::vector<ROperator*> fAddlOp;  // temporary vector to store pointer that will be moved in a unique_ptr
+
+public:
+    virtual ~RFunction_Update() {}
+    RFunction_Update() {}
+    RFunction_Update(FunctionTarget target, GraphType gType);
+
+    virtual void AddInitializedTensors(const std::vector<std::vector<std::string>>&) {};
+    virtual void Initialize() {};
+    virtual void AddLayerNormalization(int, float, size_t, const std::string&,
+                                       const std::string&, const std::string&, const std::string&) {};
+    void AddInputTensors(const std::vector<std::vector<std::size_t>>& fInputShape);
+    std::shared_ptr<RModel> GetFunctionBlock() {
+        return function_block;
+    }
+
+    std::string GenerateModel(const std::string& filename, long read_pos=0, long block_size=1);
+    std::string Generate(const std::vector<std::string>& inputPtrs);
+    FunctionTarget GetFunctionTarget() {
+        return fTarget;
+    }
 };
 
-class RFunction_Aggregate: public RFunction{
-        protected:
-                FunctionReducer fReducer;
-        public:
-        virtual ~RFunction_Aggregate(){}
-        RFunction_Aggregate(){}
-                RFunction_Aggregate(FunctionReducer reducer): fReducer(reducer){
-                        fType = FunctionType::AGGREGATE;
-                }
-        virtual std::string GenerateModel() = 0;
-        std::string GetFunctionName(){
-                return fFuncName;
-        }
-        FunctionReducer GetFunctionReducer(){
-                return fReducer;
-        }
-        std::string Generate(std::size_t num_features, const std::vector<std::string>& inputTensors){
-                std::string inferFunc = fFuncName+"("+std::to_string(num_features)+",{";
-                for(auto&it : inputTensors){
-                        inferFunc+=it;
-                        inferFunc+=",";
-                }
-                inferFunc.pop_back();
-                inferFunc+="});";
-                return inferFunc;
-        }
+class RFunction_Aggregate: public RFunction {
+protected:
+    FunctionReducer fReducer;
+public:
+    virtual ~RFunction_Aggregate() {}
+    RFunction_Aggregate() {}
+    RFunction_Aggregate(FunctionReducer reducer): fReducer(reducer) {
+        fType = FunctionType::AGGREGATE;
+    }
+    virtual std::string GenerateModel() = 0;
+    std::string GetFunctionName() {
+        return fFuncName;
+    }
+    FunctionReducer GetFunctionReducer() {
+        return fReducer;
+    }
+    std::string Generate(std::size_t num_features, const std::vector<std::string>& inputTensors);
 
 };
 
diff --git a/tmva/sofie/inc/TMVA/RFunction_MLP.hxx b/tmva/sofie/inc/TMVA/RFunction_MLP.hxx
index d396ed47b6406..c838b4c1c166c 100644
--- a/tmva/sofie/inc/TMVA/RFunction_MLP.hxx
+++ b/tmva/sofie/inc/TMVA/RFunction_MLP.hxx
@@ -1,117 +1,41 @@
 #ifndef TMVA_SOFIE_RFUNCTION_MLP
 #define TMVA_SOFIE_RFUNCTION_MLP
 
-
-#include "TMVA/SOFIE_common.hxx"
-#include "TMVA/ROperator_Concat.hxx"
-#include "TMVA/ROperator_Gemm.hxx"
-#include "TMVA/ROperator_LayerNormalization.hxx"
-#include "TMVA/ROperator_Relu.hxx"
 #include "TMVA/RFunction.hxx"
-#include "TMVA/RModel_GNN.hxx"
 
-#include <sstream>
-#include <algorithm>
-#include <iterator>
-#include <iomanip>
-#include <limits>
-#include <cassert>
-#include <memory>
 #include <vector>
 
-#include <iostream>
 
-namespace TMVA{
-namespace Experimental{
-namespace SOFIE{
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
 
 enum class Activation {
-   RELU = 0x0,
-   Invalid = 0x1,
+    RELU = 0x0,
+    Invalid = 0x1,
 };
 
-class RFunction_MLP: public RFunction_Update{
-    private:
-        Int_t fNumLayers;           // Number of Layers in MLP
-        Activation fActivationFunction;
-        bool  fActivateFinal;       // if True, fActivationFunction is applied as the activation for the last layer
-        std::vector<std::string> fKernelTensors;
-        std::vector<std::string> fBiasTensors;
-
-    public:
-        virtual ~RFunction_MLP(){}
-        RFunction_MLP(FunctionTarget target, Int_t numLayers, Activation activation_function=Activation::RELU, bool activate_final=false, GraphType gType=GraphType::GNN):
-        RFunction_Update(target, gType), fNumLayers(numLayers), fActivationFunction(activation_function), fActivateFinal(activate_final){
-            if(fActivationFunction == Activation::Invalid){
-                throw std::runtime_error("TMVA SOFIE GNN doesn't currently supports the provided activation function for " + fFuncName + " update.");
-            }
-
-            // assuming all the linear layers has a kernel and a bias initialized tensors
-            if(fActivateFinal){
-                function_block->AddOutputTensorNameList({fFuncName+"Relu"+std::to_string(fNumLayers)});
-            } else{
-                function_block->AddOutputTensorNameList({fFuncName+"Gemm"+std::to_string(fNumLayers)});
-            }
-        }
-
-        void Initialize(){
-            
-            std::string fGemmInput;
-            if(fGraphType == GraphType::GNN){            
-                std::unique_ptr<ROperator> op_concat;
-                op_concat.reset(new ROperator_Concat<float>(fInputTensors,1,0,fFuncName+"InputConcat"));
-                function_block->AddOperator(std::move(op_concat));
-                fGemmInput = fFuncName+"InputConcat";
-
-            } else if(fGraphType == GraphType::GraphIndependent){
-                fGemmInput = fInputTensors[0];
-            }
-
-            std::unique_ptr<ROperator> op_gemm;
-            for(int i=0; i<fNumLayers-1; ++i){
-                op_gemm.reset(new ROperator_Gemm<float>(1.0,1.0,0,0,fGemmInput,UTILITY::Clean_name(fKernelTensors[i]),UTILITY::Clean_name(fBiasTensors[i]),fFuncName+"Gemm"+std::to_string(i)));
-                function_block->AddOperator(std::move(op_gemm));
-                fGemmInput = fFuncName+"Gemm"+i;
-                if (fActivationFunction == Activation::RELU){
-                    std::unique_ptr<ROperator> op_relu;
-                    op_relu.reset(new ROperator_Relu<float>(fFuncName+"Gemm"+std::to_string(i), fFuncName+"Relu"+std::to_string(i)));
-                    function_block->AddOperator(std::move(op_relu));
-                    fGemmInput = fFuncName+"Relu"+i;
-
-                }       
-            }
-
-            op_gemm.reset(new ROperator_Gemm<float>(1.0,1.0,0,0,fGemmInput,UTILITY::Clean_name(fKernelTensors.back()),UTILITY::Clean_name(fBiasTensors.back()),fFuncName+"Gemm"+std::to_string(fNumLayers)));
-            function_block->AddOperator(std::move(op_gemm));
-            if(fActivateFinal){
-                if (fActivationFunction == Activation::RELU){
-                    std::unique_ptr<ROperator> op_relu;
-                    op_relu.reset(new ROperator_Relu<float>(fFuncName+"Gemm"+std::to_string(fNumLayers), fFuncName+"Relu"+std::to_string(fNumLayers)));
-                    function_block->AddOperator(std::move(op_relu));
-                } 
-            }
+class RFunction_MLP: public RFunction_Update {
+private:
+    Int_t fNumLayers;           // Number of Layers in MLP
+    Activation fActivationFunction;
+    bool  fActivateFinal;       // if True, fActivationFunction is applied as the activation for the last layer
+    std::vector<std::string> fKernelTensors;
+    std::vector<std::string> fBiasTensors;
 
+public:
+    virtual ~RFunction_MLP() {}
+    RFunction_MLP(FunctionTarget target, Int_t numLayers, Activation activation_function=Activation::RELU, bool activate_final=false, GraphType gType=GraphType::GNN);
 
-            if(fAddlOp.size()){
-                for(auto &i:fAddlOp){
-                    std::unique_ptr<ROperator> tmp(i);
-                    function_block->AddOperator(std::move(tmp));                
-                }
-            }
-        }
+    void Initialize();
 
-        void AddLayerNormalization(int axis, float epsilon, size_t stashType, const std::string &nameX,
-                                    const std::string &nameScale, const std::string &nameB, const std::string &nameY){
-            auto op_layerNorm = new ROperator_LayerNormalization<float>(axis, epsilon, stashType, nameX,
-                                                                        nameScale, nameB, nameY, "", "");
-            fAddlOp.push_back((op_layerNorm));
-        }
-        
+    void AddLayerNormalization(int axis, float epsilon, size_t stashType, const std::string &nameX,
+                               const std::string &nameScale, const std::string &nameB, const std::string &nameY);
 
-        void AddInitializedTensors(const std::vector<std::vector<std::string>>& initialized_tensors){
-                fKernelTensors = initialized_tensors[0];
-                fBiasTensors   = initialized_tensors[1];
-        }
+    void AddInitializedTensors(const std::vector<std::vector<std::string>>& initialized_tensors) {
+        fKernelTensors = initialized_tensors[0];
+        fBiasTensors   = initialized_tensors[1];
+    }
 };
 
 } // SOFIE
diff --git a/tmva/sofie/inc/TMVA/RFunction_Mean.hxx b/tmva/sofie/inc/TMVA/RFunction_Mean.hxx
index 863675a1216d4..ba9dda6cd45cc 100644
--- a/tmva/sofie/inc/TMVA/RFunction_Mean.hxx
+++ b/tmva/sofie/inc/TMVA/RFunction_Mean.hxx
@@ -2,38 +2,19 @@
 #define TMVA_SOFIE_RFUNCTION_MEAN
 
 #include "TMVA/RFunction.hxx"
-#include "TMVA/RModel_GNN.hxx"
 
-#include <sstream>
-#include <algorithm>
-#include <iterator>
-#include <iomanip>
-#include <limits>
-#include <cassert>
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
 
-namespace TMVA{
-namespace Experimental{
-namespace SOFIE{
+class RFunction_Mean: public RFunction_Aggregate {
 
-class RFunction_Mean: public RFunction_Aggregate{
-    
-    public:
-        RFunction_Mean():RFunction_Aggregate(FunctionReducer::MEAN){
-            fFuncName = "Aggregate_by_Mean";
-        }
+public:
+    RFunction_Mean():RFunction_Aggregate(FunctionReducer::MEAN) {
+        fFuncName = "Aggregate_by_Mean";
+    }
 
-        std::string GenerateModel(){
-            std::string modelGenerationString;
-            modelGenerationString = "\n//--------- GNN_Aggregate_Function---"+fFuncName+"\n";
-            modelGenerationString += "std::vector<float> "+fFuncName+"(const int& num_features, const std::vector<std::vector<float>::iterator>& inputs){\n";
-            modelGenerationString += "\tstd::vector<float> result(num_features,0);\n";
-            modelGenerationString += "\tfor(auto &it:inputs){\n";
-            modelGenerationString += "\t\tstd::transform(result.begin(), result.end(), it, result.begin(), std::plus<float>());\n\t}\n";
-            modelGenerationString += "\tfor_each(result.begin(), result.end(), [&result](float &x){ x /= result.size();\n";
-            modelGenerationString += "\treturn result;\n}";
-            return modelGenerationString;
-        }
-                
+    std::string GenerateModel();
 };
 
 } //SOFIE
diff --git a/tmva/sofie/inc/TMVA/RFunction_Sum.hxx b/tmva/sofie/inc/TMVA/RFunction_Sum.hxx
index 478ff1492410e..db7e9d5dabbf9 100644
--- a/tmva/sofie/inc/TMVA/RFunction_Sum.hxx
+++ b/tmva/sofie/inc/TMVA/RFunction_Sum.hxx
@@ -3,37 +3,19 @@
 
 
 #include "TMVA/RFunction.hxx"
-#include "TMVA/RModel_GNN.hxx"
-
-#include <sstream>
-#include <algorithm>
-#include <iterator>
-#include <iomanip>
-#include <limits>
-#include <cassert>
-
-namespace TMVA{
-namespace Experimental{
-namespace SOFIE{
-
-class RFunction_Sum: public RFunction_Aggregate{
-    
-    public:
-        RFunction_Sum():RFunction_Aggregate(FunctionReducer::SUM){
-            fFuncName = "Aggregate_by_Sum";
-        }
-
-        std::string GenerateModel(){
-            std::string modelGenerationString;
-            modelGenerationString = "\n//--------- GNN_Aggregate_Function---"+fFuncName+"\n";
-            modelGenerationString += "std::vector<float> "+fFuncName+"(const int& num_features, const std::vector<float*>& inputs){\n";
-            modelGenerationString += "\tstd::vector<float> result(num_features,0);\n";
-            modelGenerationString += "\tfor(auto &it:inputs){\n";
-            modelGenerationString += "\t\tstd::transform(result.begin(), result.end(), it, result.begin(), std::plus<float>());\n\t}\n";
-            modelGenerationString += "\treturn result;\n}";
-            return modelGenerationString;
-        }
 
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+class RFunction_Sum: public RFunction_Aggregate {
+
+public:
+    RFunction_Sum():RFunction_Aggregate(FunctionReducer::SUM) {
+        fFuncName = "Aggregate_by_Sum";
+    }
+
+    std::string GenerateModel();
 };
 
 } //SOFIE
diff --git a/tmva/sofie/inc/TMVA/RModel_Base.hxx b/tmva/sofie/inc/TMVA/RModel_Base.hxx
index 724d00324dd4c..399e06c7f24d4 100644
--- a/tmva/sofie/inc/TMVA/RModel_Base.hxx
+++ b/tmva/sofie/inc/TMVA/RModel_Base.hxx
@@ -5,7 +5,6 @@
 #include <unordered_set>
 #include <vector>
 #include <unordered_map>
-#include <iostream>
 #include <memory>
 #include <ctime>
 #include <set>
@@ -84,7 +83,7 @@ public:
     std::string ReturnGenerated() {
         return fGC;
     }
-    void OutputGenerated(std::string filename = "");
+    void OutputGenerated(std::string filename = "", bool append = false);
     void SetFilename(std::string filename) {
         fName = filename;
     }
diff --git a/tmva/sofie/inc/TMVA/ROperator_Shape.hxx b/tmva/sofie/inc/TMVA/ROperator_Shape.hxx
index 639d7727b3d06..1e0ff9b1fb62f 100644
--- a/tmva/sofie/inc/TMVA/ROperator_Shape.hxx
+++ b/tmva/sofie/inc/TMVA/ROperator_Shape.hxx
@@ -6,7 +6,6 @@
 #include "TMVA/RModel.hxx"
 
 #include <sstream>
-#include <iostream>
 #include<sstream>
 #include<vector>
 #include <iterator>
diff --git a/tmva/sofie/inc/TMVA/SOFIE_common.hxx b/tmva/sofie/inc/TMVA/SOFIE_common.hxx
index 7353d12681226..16dbb6a64903e 100644
--- a/tmva/sofie/inc/TMVA/SOFIE_common.hxx
+++ b/tmva/sofie/inc/TMVA/SOFIE_common.hxx
@@ -2,7 +2,6 @@
 #define TMVA_SOFIE_SOFIE_COMMON
 
 #include "TMVA/RTensor.hxx"
-// #include "TMVA/Types.h"
 
 #include <stdexcept>
 #include <type_traits>
diff --git a/tmva/sofie/src/RFunction.cxx b/tmva/sofie/src/RFunction.cxx
new file mode 100644
index 0000000000000..f2ec5da623db7
--- /dev/null
+++ b/tmva/sofie/src/RFunction.cxx
@@ -0,0 +1,93 @@
+#include "TMVA/RModel.hxx"
+#include "TMVA/RFunction.hxx"
+
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+
+
+RFunction_Update::RFunction_Update(FunctionTarget target, GraphType gType): fTarget(target), fGraphType(gType) {
+    switch(target) {
+    case FunctionTarget::EDGES: {
+        fFuncName = "edge_update";
+        break;
+    }
+    case FunctionTarget::NODES: {
+        fFuncName = "node_update";
+        break;
+    }
+    case FunctionTarget::GLOBALS: {
+        fFuncName = "global_update";
+        break;
+    }
+    default:
+        throw std::runtime_error("Invalid target for Update function");
+    }
+    fType = FunctionType::UPDATE;
+    function_block = std::make_unique<RModel>(fFuncName);
+
+    if(fGraphType == GraphType::GNN) {
+        if(fTarget == FunctionTarget::EDGES) {
+            fInputTensors = {"edge","receiver","sender","global"};
+        } else if(fTarget == FunctionTarget::NODES || fTarget == FunctionTarget::GLOBALS) {
+            fInputTensors = {"edge","node","global"};
+        }
+
+    } else if(fGraphType == GraphType::GraphIndependent) {
+        if(fTarget == FunctionTarget::EDGES) {
+            fInputTensors = {"edge"};
+        } else if(fTarget == FunctionTarget::NODES) {
+            fInputTensors = {"node"};
+        } else {
+            fInputTensors = {"global"};
+        }
+    }
+}
+
+void RFunction_Update::AddInputTensors(const std::vector<std::vector<std::size_t>>& fInputShape) {
+    for(long unsigned int i=0; i<fInputShape.size(); ++i) {
+        function_block->AddInputTensorInfo(fInputTensors[i],ETensorType::FLOAT, fInputShape[i]);
+        function_block->AddInputTensorName(fInputTensors[i]);
+    }
+}
+
+std::string RFunction_Update::GenerateModel(const std::string& filename, long read_pos, long block_size) {
+    function_block->SetFilename(filename);
+    // use batch size as block size in RModel::generate
+    function_block->Generate(Options::kGNNComponent,block_size,read_pos);
+    std::string modelGenerationString;
+    modelGenerationString = "\n//--------- GNN_Update_Function---"+fFuncName+"\n"+function_block->ReturnGenerated();
+    return modelGenerationString;
+}
+
+std::string RFunction_Update::Generate(const std::vector<std::string>& inputPtrs) {
+    std::string inferFunc = fFuncName+".infer(";
+    for(auto&it : inputPtrs) {
+        inferFunc+=it;
+        inferFunc+=",";
+    }
+    inferFunc.pop_back();
+    inferFunc+=");";
+    return inferFunc;
+}
+
+
+std::string RFunction_Aggregate::Generate(std::size_t num_features, const std::vector<std::string>& inputTensors) {
+    std::string inferFunc = fFuncName+"("+std::to_string(num_features)+",{";
+    for(auto&it : inputTensors) {
+        inferFunc+=it;
+        inferFunc+=",";
+    }
+    inferFunc.pop_back();
+    inferFunc+="});";
+    return inferFunc;
+}
+
+
+
+
+}
+}
+}
diff --git a/tmva/sofie/src/RFunction_MLP.cxx b/tmva/sofie/src/RFunction_MLP.cxx
new file mode 100644
index 0000000000000..cd2ce2fbb929e
--- /dev/null
+++ b/tmva/sofie/src/RFunction_MLP.cxx
@@ -0,0 +1,83 @@
+#include "TMVA/RFunction_MLP.hxx"
+
+
+#include "TMVA/SOFIE_common.hxx"
+#include "TMVA/ROperator_Concat.hxx"
+#include "TMVA/ROperator_Gemm.hxx"
+#include "TMVA/ROperator_LayerNormalization.hxx"
+#include "TMVA/ROperator_Relu.hxx"
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+RFunction_MLP::RFunction_MLP(FunctionTarget target, Int_t numLayers, Activation activation_function, bool activate_final, GraphType gType):
+    RFunction_Update(target, gType), fNumLayers(numLayers), fActivationFunction(activation_function), fActivateFinal(activate_final) {
+    if(fActivationFunction == Activation::Invalid) {
+        throw std::runtime_error("TMVA SOFIE GNN doesn't currently supports the provided activation function for " + fFuncName + " update.");
+    }
+
+    // assuming all the linear layers has a kernel and a bias initialized tensors
+    if(fActivateFinal) {
+        function_block->AddOutputTensorNameList({fFuncName+"Relu"+std::to_string(fNumLayers)});
+    } else {
+        function_block->AddOutputTensorNameList({fFuncName+"Gemm"+std::to_string(fNumLayers)});
+    }
+}
+
+void RFunction_MLP::Initialize() {
+
+    std::string fGemmInput;
+    if(fGraphType == GraphType::GNN) {
+        std::unique_ptr<ROperator> op_concat;
+        op_concat.reset(new ROperator_Concat<float>(fInputTensors,1,0,fFuncName+"InputConcat"));
+        function_block->AddOperator(std::move(op_concat));
+        fGemmInput = fFuncName+"InputConcat";
+
+    } else if(fGraphType == GraphType::GraphIndependent) {
+        fGemmInput = fInputTensors[0];
+    }
+
+    std::unique_ptr<ROperator> op_gemm;
+    for(int i=0; i<fNumLayers-1; ++i) {
+        op_gemm.reset(new ROperator_Gemm<float>(1.0,1.0,0,0,fGemmInput,UTILITY::Clean_name(fKernelTensors[i]),UTILITY::Clean_name(fBiasTensors[i]),fFuncName+"Gemm"+std::to_string(i)));
+        function_block->AddOperator(std::move(op_gemm));
+        fGemmInput = fFuncName+"Gemm"+i;
+        if (fActivationFunction == Activation::RELU) {
+            std::unique_ptr<ROperator> op_relu;
+            op_relu.reset(new ROperator_Relu<float>(fFuncName+"Gemm"+std::to_string(i), fFuncName+"Relu"+std::to_string(i)));
+            function_block->AddOperator(std::move(op_relu));
+            fGemmInput = fFuncName+"Relu"+i;
+
+        }
+    }
+
+    op_gemm.reset(new ROperator_Gemm<float>(1.0,1.0,0,0,fGemmInput,UTILITY::Clean_name(fKernelTensors.back()),UTILITY::Clean_name(fBiasTensors.back()),fFuncName+"Gemm"+std::to_string(fNumLayers)));
+    function_block->AddOperator(std::move(op_gemm));
+    if(fActivateFinal) {
+        if (fActivationFunction == Activation::RELU) {
+            std::unique_ptr<ROperator> op_relu;
+            op_relu.reset(new ROperator_Relu<float>(fFuncName+"Gemm"+std::to_string(fNumLayers), fFuncName+"Relu"+std::to_string(fNumLayers)));
+            function_block->AddOperator(std::move(op_relu));
+        }
+    }
+
+
+    if(fAddlOp.size()) {
+        for(auto &i:fAddlOp) {
+            std::unique_ptr<ROperator> tmp(i);
+            function_block->AddOperator(std::move(tmp));
+        }
+    }
+}
+
+void RFunction_MLP::AddLayerNormalization(int axis, float epsilon, size_t stashType, const std::string &nameX,
+        const std::string &nameScale, const std::string &nameB, const std::string &nameY) {
+    auto op_layerNorm = new ROperator_LayerNormalization<float>(axis, epsilon, stashType, nameX,
+            nameScale, nameB, nameY, "", "");
+    fAddlOp.push_back((op_layerNorm));
+}
+
+}
+}
+}
diff --git a/tmva/sofie/src/RFunction_Mean.cxx b/tmva/sofie/src/RFunction_Mean.cxx
new file mode 100644
index 0000000000000..554d782d69b3a
--- /dev/null
+++ b/tmva/sofie/src/RFunction_Mean.cxx
@@ -0,0 +1,22 @@
+#include "TMVA/RFunction_Mean.hxx"
+
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+std::string RFunction_Mean::GenerateModel() {
+    std::string modelGenerationString;
+    modelGenerationString = "\n//--------- GNN_Aggregate_Function---"+fFuncName+"\n";
+    modelGenerationString += "std::vector<float> "+fFuncName+"(const int& num_features, const std::vector<std::vector<float>::iterator>& inputs){\n";
+    modelGenerationString += "\tstd::vector<float> result(num_features,0);\n";
+    modelGenerationString += "\tfor(auto &it:inputs){\n";
+    modelGenerationString += "\t\tstd::transform(result.begin(), result.end(), it, result.begin(), std::plus<float>());\n\t}\n";
+    modelGenerationString += "\tfor_each(result.begin(), result.end(), [&result](float &x){ x /= result.size();\n";
+    modelGenerationString += "\treturn result;\n}";
+    return modelGenerationString;
+}
+
+}
+}
+}
diff --git a/tmva/sofie/src/RFunction_Sum.cxx b/tmva/sofie/src/RFunction_Sum.cxx
new file mode 100644
index 0000000000000..7db9d9a2f9f7a
--- /dev/null
+++ b/tmva/sofie/src/RFunction_Sum.cxx
@@ -0,0 +1,21 @@
+#include "TMVA/RFunction_Sum.hxx"
+
+
+namespace TMVA {
+namespace Experimental {
+namespace SOFIE {
+
+std::string RFunction_Sum::GenerateModel() {
+    std::string modelGenerationString;
+    modelGenerationString = "\n//--------- GNN_Aggregate_Function---"+fFuncName+"\n";
+    modelGenerationString += "std::vector<float> "+fFuncName+"(const int& num_features, const std::vector<float*>& inputs){\n";
+    modelGenerationString += "\tstd::vector<float> result(num_features,0);\n";
+    modelGenerationString += "\tfor(auto &it:inputs){\n";
+    modelGenerationString += "\t\tstd::transform(result.begin(), result.end(), it, result.begin(), std::plus<float>());\n\t}\n";
+    modelGenerationString += "\treturn result;\n}";
+    return modelGenerationString;
+}
+
+}
+}
+}
diff --git a/tmva/sofie/src/RModel.cxx b/tmva/sofie/src/RModel.cxx
index 20d42d873a683..c55fedc2561f3 100644
--- a/tmva/sofie/src/RModel.cxx
+++ b/tmva/sofie/src/RModel.cxx
@@ -536,70 +536,68 @@ void RModel::ReadInitializedTensorsFromFile(long pos) {
 }
 
 long RModel::WriteInitializedTensorsToFile(std::string filename) {
-   // Determine the file extension based on the weight file type
-   std::string fileExtension;
-   switch (fWeightFile) {
-      case WeightFileType::None:
-         fileExtension = ".dat";
-         break;
-      case WeightFileType::RootBinary:
-         fileExtension = ".root";
-         break;
-      case WeightFileType::Text:
-         fileExtension = ".dat";
-         break;
-   }
-
-   // If filename is empty, use the model name as the base filename
-   if (filename.empty()) {
-      filename = fFileName + fileExtension;
-   }
-
-   // Write the initialized tensors to the file
-   if (fWeightFile == WeightFileType::RootBinary) {
+    // Determine the file extension based on the weight file type
+    std::string fileExtension;
+    switch (fWeightFile) {
+    case WeightFileType::None:
+        fileExtension = ".dat";
+        break;
+    case WeightFileType::RootBinary:
+        fileExtension = ".root";
+        break;
+    case WeightFileType::Text:
+        fileExtension = ".dat";
+        break;
+    }
+
+    // If filename is empty, use the model name as the base filename
+    if (filename.empty()) {
+        filename = fFileName + fileExtension;
+    }
+
+    // Write the initialized tensors to the file
+    if (fWeightFile == WeightFileType::RootBinary) {
         if(fIsGNNComponent || fIsGNN) {
-            throw std::runtime_error("SOFIE-GNN yet not supports writing to a ROOT file.")
-        }
-      std::unique_ptr<TFile> outputFile(TFile::Open(filename.c_str(), "UPDATE"));
-
-      std::string dirName = fName + "_weights";
-      // check if directory exists, in case delete to replace with new one
-      if (outputFile->GetKey(dirName.c_str()))
-         outputFile->rmdir(dirName.c_str());
-
-      auto outputDir = outputFile->mkdir(dirName.c_str());
-
-      for (const auto& item : fInitializedTensors) {
-         std::string tensorName = "tensor_" + item.first;
-         size_t length = 1;
-         length = ConvertShapeToLength(item.second.fShape);
-         if(item.second.fType == ETensorType::FLOAT){
-            const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
-            const float* data = (std::static_pointer_cast<float>(item.second.fData)).get();
-            std::vector<float> tensorDataVector(data , data + length);
-            outputDir->WriteObjectAny(&tensorDataVector, "std::vector<float>", tensorName.c_str());
-         }
-         else if(item.second.fType == ETensorType::DOUBLE){
-            const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
-            const double* data = (std::static_pointer_cast<double>(item.second.fData)).get();
-            std::vector<double> tensorDataVector(data , data + length);
-            outputDir->WriteObjectAny(&tensorDataVector, "std::vector<double>", tensorName.c_str());
-         }
-         else if(item.second.fType == ETensorType::INT64) {
-            const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
-            const int64_t* data = (std::static_pointer_cast<int64_t>(item.second.fData)).get();
-            std::vector<int64_t> tensorDataVector(data , data + length);
-            outputDir->WriteObjectAny(&tensorDataVector, "std::vector<int64_t>", tensorName.c_str());
-         }
-      }
-      outputFile->Write(filename.c_str());
-    
-        // this needs to be changed, similar to the text file
-        return 0;
-   }
+            throw std::runtime_error("SOFIE-GNN yet not supports writing to a ROOT file.");
+        }
+        std::unique_ptr<TFile> outputFile(TFile::Open(filename.c_str(), "UPDATE"));
+
+        std::string dirName = fName + "_weights";
+        // check if directory exists, in case delete to replace with new one
+        if (outputFile->GetKey(dirName.c_str()))
+            outputFile->rmdir(dirName.c_str());
+
+        auto outputDir = outputFile->mkdir(dirName.c_str());
+
+        for (const auto& item : fInitializedTensors) {
+            std::string tensorName = "tensor_" + item.first;
+            size_t length = 1;
+            length = ConvertShapeToLength(item.second.fShape);
+            if(item.second.fType == ETensorType::FLOAT) {
+                const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
+                const float* data = (std::static_pointer_cast<float>(item.second.fData)).get();
+                std::vector<float> tensorDataVector(data, data + length);
+                outputDir->WriteObjectAny(&tensorDataVector, "std::vector<float>", tensorName.c_str());
+            }
+            else if(item.second.fType == ETensorType::DOUBLE) {
+                const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
+                const double* data = (std::static_pointer_cast<double>(item.second.fData)).get();
+                std::vector<double> tensorDataVector(data, data + length);
+                outputDir->WriteObjectAny(&tensorDataVector, "std::vector<double>", tensorName.c_str());
+            }
+            else if(item.second.fType == ETensorType::INT64) {
+                const std::shared_ptr<void> ptr = item.second.fData; // shared_ptr<void> instance
+                const int64_t* data = (std::static_pointer_cast<int64_t>(item.second.fData)).get();
+                std::vector<int64_t> tensorDataVector(data, data + length);
+                outputDir->WriteObjectAny(&tensorDataVector, "std::vector<int64_t>", tensorName.c_str());
+            }
+        }
+        outputFile->Write(filename.c_str());
 
-   // Write the initialized tensors to a text file
-   if (fWeightFile == WeightFileType::Text) {
+        // this needs to be changed, similar to the text file
+        return -1;
+   
+    } else if (fWeightFile == WeightFileType::Text) {
         std::ofstream f;
         if(fIsGNNComponent) {
             // appending all GNN components into the same file
@@ -629,6 +627,8 @@ long RModel::WriteInitializedTensorsToFile(std::string filename) {
         long curr_pos = f.tellp();
         f.close();
         return curr_pos;
+    } else {
+        return -1;
     }
 }
 
@@ -714,16 +714,12 @@ void RModel::HeadInitializedTensors(std::string name, int n_print) {
     }
 
     std::cout << "data: [" << std::endl;
-    //switch(it->second.type){
-    //   case ETensorType::FLOAT : {
     if (it->second.fType == ETensorType::FLOAT) {
         auto converted_data = std::static_pointer_cast<float>(it->second.fData).get();
         for (int i =0; i < n_print; i++) {
             std::cout << converted_data[i];
             if (i < n_print - 1) std::cout << " ,";
         }
-        //   break;
-        // }
     }
     if (ellipsis) std::cout << ", ...";
     std::cout << "]" << std::endl;
@@ -731,7 +727,7 @@ void RModel::HeadInitializedTensors(std::string name, int n_print) {
 }
 
 void RModel::OutputGenerated(std::string filename, bool append) {
-    
+
     RModel_Base::OutputGenerated(filename, append);
 
     // write weights in a text file
diff --git a/tmva/sofie/src/RModel_GNN.cxx b/tmva/sofie/src/RModel_GNN.cxx
index ba5545f750850..0a45292f8bf73 100644
--- a/tmva/sofie/src/RModel_GNN.cxx
+++ b/tmva/sofie/src/RModel_GNN.cxx
@@ -3,7 +3,6 @@
 #include <fstream>
 #include <limits>
 
-#include <iostream>
 #include "TMVA/RModel_GNN.hxx"
 #include "TMVA/RFunction.hxx"
 
@@ -129,10 +128,10 @@ void RModel_GNN::Generate() {
     next_pos = globals_update_block->GetFunctionBlock()->WriteInitializedTensorsToFile(fName+".dat");
     fGC+="};\n}\n";
 
-    // correct for difference in global size
+    // correct for difference in global size  (check shape[1] of output og globals update)
     auto num_global_features_input = num_global_features;
-    if(globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[0] != num_global_features) {
-        num_global_features = globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[0];
+    if(globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1] != num_global_features) {
+        num_global_features = globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1];
     }
 
     fGC+=edge_node_agg_block->GenerateModel();
@@ -225,10 +224,21 @@ void RModel_GNN::Generate() {
            ", input_graph.node_data.GetData() + (k + 1) * " + n_size_input +
            ", fNodeInputs.begin() + k * " + n_size_input + ");\n";
     fGC += "}\n";
-    std::vector<std::string> Node_Edge_Aggregate_String;
+    // reset initial aggregate edge vector to zero
+    fGC += "\nstd::fill(fNodeEdgeAggregate.begin(), fNodeEdgeAggregate.end(), 0.);\n";
+    // fGlobInputs is size { nedges, ngloblas}. It needs to be here { nnodes, nglobals}
+    // if number of nodes is larger than edges we need to resize it and copy values
+    if (num_nodes > num_edges) {
+        fGC += "\n// resize global vector feature to number of nodes\n";
+        fGC += "fGlobInputs.resize( " + std::to_string(num_nodes * num_global_features_input) + ");";
+        fGC += "for (size_t k = " + std::to_string(num_edges) + "; k < " + std::to_string(num_nodes) + "; k++)";
+            fGC += "   std::copy(fGlobInputs.begin(), fGlobInputs.begin() + " + std::to_string(num_global_features_input) +
+                   " , fGlobInputs.begin() + k * " + std::to_string(num_global_features_input) + ");\n";
+    }
 
     // aggregating edge if it's a receiver node and then updating corresponding node
     for(int i=0; i<num_nodes; ++i) {
+        std::vector<std::string> Node_Edge_Aggregate_String;
         for(int k=0; k<num_edges; ++k) {
             if(receivers[k] == i) {
                 Node_Edge_Aggregate_String.emplace_back("input_graph.edge_data.GetData()+"+std::to_string(k*num_edge_features));
@@ -236,12 +246,11 @@ void RModel_GNN::Generate() {
         }
 
         // when node is not a receiver, fill the aggregated vector with 0 values
-        if(Node_Edge_Aggregate_String.size()==0) {
-            fGC+="fNodeEdgeAggregate.insert(fNodeEdgeAggregate.end(), "+std::to_string(num_edge_features)+", 0);";
-        } else {
+        if(Node_Edge_Aggregate_String.size()!=0) {
             fGC+="\nfNodeAggregateTemp = ";
             fGC+=edge_node_agg_block->Generate(num_edge_features, {Node_Edge_Aggregate_String});                     // aggregating edge attributes per node
-            fGC+="\nfNodeEdgeAggregate.insert(fNodeEdgeAggregate.begin(), fNodeAggregateTemp.begin(), fNodeAggregateTemp.end());";
+            fGC += "\nstd::copy(fNodeAggregateTemp.begin(), fNodeAggregateTemp.end(), fNodeEdgeAggregate.begin() + " +
+                   std::to_string(num_edge_features * i) + ");";
         }
     }
 
@@ -249,7 +258,6 @@ void RModel_GNN::Generate() {
     fGC+="fNodeUpdates = ";
     fGC+=nodes_update_block->Generate({"fNodeEdgeAggregate.data()","fNodeInputs.data()","fGlobInputs.data()"});    // computing updated node attributes
     fGC+="\n";
-    Node_Edge_Aggregate_String.clear();
 
     if(num_node_features != num_node_features_input) {
         fGC += "\n//  resize node graph data since output feature size is not equal to input size\n";
@@ -285,9 +293,9 @@ void RModel_GNN::Generate() {
     // computing updated global attributes
     fGC += "std::vector<float> Global_Data = ";
     fGC += globals_update_block->Generate({"Edge_Global_Aggregate.data()","Node_Global_Aggregate.data()", "input_graph.global_data.GetData()"});
-    if(globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1] != num_global_features) {
-        num_global_features = globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1];
-        fGC+="\ninput_graph.global_data = input_graph.global_data.Resize({"+std::to_string(num_global_features)+"});";
+    if(num_global_features != num_global_features_input) {
+        fGC += "\n//  resize global graph data since output feature size is not equal to input size\n";
+        fGC+="input_graph.global_data = input_graph.global_data.Resize({"+std::to_string(num_global_features)+"});\n";
     }
     fGC += "\nstd::copy(Global_Data.begin(), Global_Data.end(), input_graph.global_data.GetData());";
     fGC+="\n}\n";
diff --git a/tmva/sofie/src/RModel_GraphIndependent.cxx b/tmva/sofie/src/RModel_GraphIndependent.cxx
index 0286d43e2ce31..bdb713542dfca 100644
--- a/tmva/sofie/src/RModel_GraphIndependent.cxx
+++ b/tmva/sofie/src/RModel_GraphIndependent.cxx
@@ -106,6 +106,14 @@ void RModel_GraphIndependent::Generate() {
     next_pos = globals_update_block->GetFunctionBlock()->WriteInitializedTensorsToFile(fName+".dat");
     fGC+="};\n}\n";
 
+    // we need to correct the output number of global features
+    auto num_global_features_input = num_global_features;
+    // global features are in shape[1]
+    if(globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1] != num_global_features) {
+        num_global_features = globals_update_block->GetFunctionBlock()->GetTensorShape(globals_update_block->GetFunctionBlock()->GetOutputTensorNames()[0])[1];
+    }
+
+
     // computing inplace on input graph
     fGC += "struct Session {\n";
     fGC += "\n// Instantiating session objects for graph components\n";
@@ -176,6 +184,13 @@ void RModel_GraphIndependent::Generate() {
     fGC += "\n// --- Global Update ---\n";
     fGC += "std::vector<float> Global_Data = ";
     fGC += globals_update_block->Generate({"input_graph.global_data.GetData()"});
+    fGC += "\n";
+
+    if(num_global_features != num_global_features_input) {
+        fGC += "\n//  resize global graph data since output feature size is not equal to input size\n";
+        fGC+="input_graph.global_data = input_graph.global_data.Resize({"+std::to_string(num_global_features)+"});\n";
+    }
+
     fGC += "\nstd::copy(Global_Data.begin(), Global_Data.end(), input_graph.global_data.GetData());";
     fGC += "\n";
 
diff --git a/tmva/sofie/src/SOFIE_common.cxx b/tmva/sofie/src/SOFIE_common.cxx
index 38b0996b70650..58e99dd43a476 100644
--- a/tmva/sofie/src/SOFIE_common.cxx
+++ b/tmva/sofie/src/SOFIE_common.cxx
@@ -1,7 +1,6 @@
 #include "TMVA/SOFIE_common.hxx"
 #include<cctype>
 #include <sstream>
-#include <iostream>
 #include <stdexcept>
 
 namespace TMVA{
diff --git a/tmva/sofie/test/EmitFromONNX.cxx.in b/tmva/sofie/test/EmitFromONNX.cxx.in
index 594586795c4a6..79ba3e7eef8fd 100644
--- a/tmva/sofie/test/EmitFromONNX.cxx.in
+++ b/tmva/sofie/test/EmitFromONNX.cxx.in
@@ -5,8 +5,6 @@
 //    This program is automatically run when the target 'TestCustomModelsFromONNX' is built.
 //    Usage example: $./sofiec indir/mymodel.onnx outdir/myname.hxx
 
-#include <iostream>
-
 #include "TMVA/RModel_Base.hxx"
 #include "TMVA/RModel.hxx"
 #include "TMVA/RModelParser_ONNX.hxx"
diff --git a/tmva/sofie/test/EmitFromRoot.cxx.in b/tmva/sofie/test/EmitFromRoot.cxx.in
index 5f9aeb786207f..9689485da92cd 100644
--- a/tmva/sofie/test/EmitFromRoot.cxx.in
+++ b/tmva/sofie/test/EmitFromRoot.cxx.in
@@ -6,8 +6,6 @@
 //           generates the required .hxx file after reading a written
 //           ROOT file which stores the object of the RModel class.
 
-#include <iostream>
-
 #include "TMVA/RModel.hxx"
 #include "TMVA/RModelParser_ONNX.hxx"
 #include "TFile.h"
diff --git a/tmva/sofie/test/GNN/EmitGNN.cxx b/tmva/sofie/test/GNN/EmitGNN.cxx
index 96070ac927432..ef81afbd193c0 100644
--- a/tmva/sofie/test/GNN/EmitGNN.cxx
+++ b/tmva/sofie/test/GNN/EmitGNN.cxx
@@ -2,8 +2,6 @@
 // Description:
 //    This program generates a RModel_GNN for testing
 
-#include <iostream>
-
 #include "TMVA/RModel_GNN.hxx"
 #include "TMVA/FunctionList.hxx"
 #include "TMVA/SOFIE_common.hxx"
diff --git a/tmva/sofie/test/GNN/EmitGraphIndependent.cxx b/tmva/sofie/test/GNN/EmitGraphIndependent.cxx
index 04ab9abd34bc7..8467a29cfb8eb 100644
--- a/tmva/sofie/test/GNN/EmitGraphIndependent.cxx
+++ b/tmva/sofie/test/GNN/EmitGraphIndependent.cxx
@@ -2,8 +2,6 @@
 // Description:
 //    This program generates a RModel_GraphIndependent for testing
 
-#include <iostream>
-
 #include "TMVA/RModel_GraphIndependent.hxx"
 #include "TMVA/FunctionList.hxx"
 #include "TMVA/SOFIE_common.hxx"
diff --git a/tmva/tmva/test/rtensor.cxx b/tmva/tmva/test/rtensor.cxx
index 6982feed50613..e751caea2233e 100644
--- a/tmva/tmva/test/rtensor.cxx
+++ b/tmva/tmva/test/rtensor.cxx
@@ -1,6 +1,6 @@
 #include <gtest/gtest.h>
 #include <TMVA/RTensor.hxx>
-#include <iostream>
+
 using namespace TMVA::Experimental;
 
 TEST(RTensor, GetElement)