Repository: bfGraph/STGraph (language: python)

Issue:
Merging Benchmark/pcsr into Benchmark/gpma



Discussion / hints:

Fix patch:
diff --git a/assets/with mod.png b/assets/with mod.png
new file mode 100644
index 0000000..b9d6547
Binary files /dev/null and b/assets/with mod.png differ
diff --git a/benchmarking/dataset/preprocess_temporal_data.py b/benchmarking/dataset/preprocess_temporal_data.py
new file mode 100644
index 0000000..54f3201
--- /dev/null
+++ b/benchmarking/dataset/preprocess_temporal_data.py
@@ -0,0 +1,128 @@
+import json
+import random
+
+
+def parse_txt_lines(lines, num_nodes):
+	id_to_pid_map = {}   # ID to processed ID
+	node_counter = 0
+
+	node_set = set()
+	edges = []
+	for line in lines:
+		parsed_line = line.split(' ')
+
+		src = int(parsed_line[0])
+		dst = int(parsed_line[1])
+
+		if src not in id_to_pid_map:
+			id_to_pid_map[src] = node_counter
+			node_counter += 1
+
+		if dst not in id_to_pid_map:
+			id_to_pid_map[dst] = node_counter
+			node_counter += 1
+
+		edges.append((id_to_pid_map[src], id_to_pid_map[dst]))
+		node_set.add(id_to_pid_map[src])
+		node_set.add(id_to_pid_map[dst])
+
+	expected_node_set = set([i for i in range(num_nodes)])
+	assert node_set == expected_node_set, "Node labelling is not continuous"
+	# print(node_set.difference(expected_node_set))
+	# quit()
+	return edges
+
+
+def preprocess_graph(edges, num_nodes):
+	# base = 2000000
+	# add_delta = 50000
+	# delete_delta = 50000
+
+	base = 100000
+	add_delta = 10000
+	delete_delta = 10000
+
+	graphs = {}
+
+	graphs["0"] = {
+					"add": list(set(edges[:base])),
+					"delete": [],
+					"neg": []
+				}
+
+	add_start_ptr = base
+	delete_start_ptr = 0
+	time = 1
+
+	while (add_start_ptr + add_delta) < len(edges):
+		graph_t = set(edges[delete_start_ptr + delete_delta:add_start_ptr + add_delta])
+		common_with_prev_graph = graph_t.intersection(set(edges[delete_start_ptr:add_start_ptr]))
+
+		# No need to add an edge thats already there in the potion carried over from previous snapshot
+		add_edges = list(set(edges[add_start_ptr:add_start_ptr + add_delta]).difference(common_with_prev_graph))
+
+		# Difference to prevent the deletion of an edge that is present in the snapshot
+		del_edges = list(set(edges[delete_start_ptr:delete_start_ptr + delete_delta]).difference(graph_t))
+
+		neg_edge_set_t = set()
+		while len(neg_edge_set_t) != len(add_edges):
+			candidate = (random.randint(0, num_nodes-1), random.randint(0, num_nodes-1))
+			if candidate not in graph_t:
+				neg_edge_set_t.add(candidate)
+
+		graphs[str(time)] = {
+								"add": add_edges,
+								"delete": del_edges,
+								"neg": list(neg_edge_set_t)
+							}
+		
+		assert len(set(edges[delete_start_ptr:add_start_ptr])) + len(add_edges) - len(del_edges) == len(graph_t)
+
+		add_start_ptr += add_delta
+		delete_start_ptr += delete_delta
+		time = time + 1
+
+	if add_start_ptr < len(edges):
+		graph_t = set(edges[delete_start_ptr+delete_delta:])
+		common_with_prev_graph = graph_t.intersection(set(edges[delete_start_ptr:add_start_ptr]))
+
+		# No need to add an edge thats already there in the potion carried over from previous snapshot
+		add_edges = list(set(edges[add_start_ptr:]).difference(common_with_prev_graph))
+
+		# Difference to prevent the deletion of an edge that is going to be added
+		del_edges = list(set(edges[delete_start_ptr:delete_start_ptr + delete_delta]).difference(graph_t))
+
+		neg_edge_set_t = set()
+		while len(neg_edge_set_t) != len(add_edges):
+			candidate = (random.randint(0, num_nodes-1), random.randint(0, num_nodes-1))
+			if candidate not in graph_t:
+				neg_edge_set_t.add(candidate)
+
+		graphs[str(time)] = {
+								"add": add_edges,
+								"delete": del_edges,
+								"neg": list(neg_edge_set_t)
+							}
+		
+		assert len(set(edges[delete_start_ptr:add_start_ptr])) + len(add_edges) - len(del_edges) == len(graph_t)
+
+		time = time+1
+
+	graph_json = {
+		"edge_mapping": {"edge_index": graphs},
+		"time_periods": time,
+	}
+
+	return graph_json
+
+
+file1 = open('sx-mathoverflow.txt', 'r')
+lines = file1.readlines()
+
+num_nodes = 24818
+edges = parse_txt_lines(lines, num_nodes)
+
+graph_json = preprocess_graph(edges, num_nodes)
+out_file = open("sx-mathoverflow-data.json", "w")
+json.dump(graph_json, out_file)
+out_file.close()
diff --git a/benchmarking/dataset_builder/dataset_builder_modified.py b/benchmarking/dataset_builder/dataset_builder_modified.py
new file mode 100644
index 0000000..45c9969
--- /dev/null
+++ b/benchmarking/dataset_builder/dataset_builder_modified.py
@@ -0,0 +1,209 @@
+"""
+    Dataset Builder for Soorah - Random Sparse Graphs
+"""
+
+import json
+import os
+from random import randint, shuffle, getrandbits
+from argparse import ArgumentParser
+import time
+from tqdm import tqdm
+import copy
+
+from rich import inspect
+
+parser = ArgumentParser(
+    description="Create custom Super Large Sparse Graph Dynamic Dataset - Soorah"
+)
+
+parser.add_argument(
+    'dataset_name',
+    help="Name of the dynamic graph dataset",
+    metavar="dataset_name"
+)
+
+parser.add_argument(
+    '-N',
+    help="Number of nodes to be present in the graph",
+    default=500,
+    metavar="num_nodes",
+    type=int
+)
+
+parser.add_argument(
+    '-M',
+    help="Muliplier that is used to get number of edges from total possible edges",
+    default=0.2,
+    metavar="edge_multiplier",
+    type=float
+)
+
+parser.add_argument(
+    '-A',
+    help="Coefficient to find number of edges to be added to the current graph",
+    default=0.1,
+    metavar="add_coeff",
+    type=float
+)
+
+parser.add_argument(
+    '-D',
+    help="Coefficient to find number of edges to be deleted from the current graph",
+    default=0.05,
+    metavar="del_coeff",
+    type=float
+)
+
+parser.add_argument(
+    '-T',
+    help="Total time stamps to be in the dynamic graph dataset",
+    default=100,
+    metavar="total_timestamps",
+    type=int
+)
+
+parser.add_argument(
+    '-L',
+    help="Low-limit multiplier for number of edges",
+    default=0.9,
+    metavar="low_limit",
+    type=float
+)
+
+parser.add_argument(
+    '-U',
+    help="Upper-limit multiplier for number of edges",
+    default=1.1,
+    metavar="upp_limit",
+    type=float
+)
+
+parser.add_argument(
+    '-F',
+    help="Max Feature size for each node of the graph",
+    default=8,
+    metavar="max_feat_size",
+    type=int
+)
+
+parser.add_argument(
+    '-I',
+    help="Increment feature size",
+    default=8,
+    metavar="inc_feat_size",
+    type=int
+)
+
+args = parser.parse_args()
+
+def create_graph(
+    num_nodes=500,
+    edge_multiplier=0.2,
+    add_coeff=0.1,
+    del_coeff=0.05,
+    total_time=100,
+):
+    graph_json = {
+        "edge_mapping": {"edge_index": {}, "edge_weight": {}},
+        "time_periods": total_time,
+        "y": [],
+    }
+
+    # edge_count = (num_nodes * (num_nodes - 1)) * edge_multiplier
+
+    # for sparse graphs edge_count is within a constant multiple of edge_multiplier
+    edge_count = num_nodes * edge_multiplier
+
+
+    # num_edges_low_limit = low_limit * edge_count
+    # num_edges_upp_limit = upp_limit * edge_count
+
+    # creating the base graph
+    edge_list = set()
+    edge_weight_list = []
+    feature_list = []
+
+    # this will have the number of edges in the graph
+    # for the current time stamp
+    # curr_time_edge_count = randint(int(num_edges_low_limit), int(num_edges_upp_limit))
+    curr_time_edge_count = edge_count
+
+    while len(edge_list) != curr_time_edge_count:
+        edge = (randint(0, num_nodes - 1), randint(0, num_nodes - 1))
+        edge_list.add(edge)
+
+    edge_weight_list = [randint(1, 1000) for i in range(len(edge_list))]
+            
+    # getting the node features for each node of the base graph
+    # a0 = time.time()
+    # feature_list = [[randint(0, 100) for _ in range(feat_size)] for _ in range(num_nodes)]
+    # a1 = time.time()
+
+    # print(f'Time to get all features: {a1-a0}')
+    
+    edge_list = list(edge_list)
+    
+    graph_json["edge_mapping"]["edge_index"]["0"] = edge_list
+    graph_json["edge_mapping"]["edge_weight"]["0"] = edge_weight_list
+    # graph_json["y"].append(feature_list)
+
+    # creating the next time stamp info from previous
+    # time stamps information
+    prev_edge_list = edge_list
+    for time_stamp in tqdm(range(1, total_time)):
+        shuffle(prev_edge_list)
+
+        del_edge_count = int(len(prev_edge_list) * del_coeff)
+        add_edge_count = int(len(prev_edge_list) * add_coeff)
+
+        curr_edge_list = set(prev_edge_list[del_edge_count:])
+        curr_edge_weight_list = [randint(1, 1000) for i in range(len(curr_edge_list))]
+        curr_time_edge_count = len(curr_edge_list) + add_edge_count
+
+        while len(curr_edge_list) != curr_time_edge_count:
+            edge = (randint(0, num_nodes - 1), randint(0, num_nodes - 1))
+            curr_edge_list.add(edge)
+
+        # adding self loops
+        # for i in range(num_nodes):
+        #     curr_edge_list.add((i,i))
+
+        curr_edge_weight_list = [randint(1, 1000) for i in range(len(curr_edge_list))]
+        # curr_feature_list = [[randint(0, 100) for _ in range(feat_size)] for _ in range(num_nodes)]
+        curr_edge_list = list(curr_edge_list)
+
+        graph_json["edge_mapping"]["edge_index"][str(time_stamp)] = curr_edge_list
+        graph_json["edge_mapping"]["edge_weight"][str(time_stamp)] = curr_edge_weight_list
+        # graph_json["y"].append(curr_feature_list)
+
+        prev_edge_list = curr_edge_list
+
+    return graph_json
+
+graph_name = args.dataset_name
+dataset_path = "../dataset/"
+
+t0 = time.time()
+
+graph_json = create_graph(args.N, args.M, args.A, args.D, args.T)
+
+folder_path = os.path.join(dataset_path, graph_name)
+
+if not os.path.exists(dataset_path + graph_name):
+    os.mkdir(folder_path)
+
+feat_size = args.I
+while feat_size <= args.F:
+    curr_graph_json = copy.deepcopy(graph_json)
+    for _ in range(args.T):
+        curr_feature_list = [[randint(0, 100) for _ in range(feat_size)] for _ in range(args.N)]
+        curr_graph_json["y"].append(curr_feature_list)
+
+    with open(f"{folder_path}/{graph_name}_{feat_size}.json", "w") as fp:
+        json.dump(curr_graph_json, fp)
+    
+    feat_size = feat_size + args.I
+
+t1 = time.time()
+
+print(f'Total time taken: {t1-t0}')
\ No newline at end of file
diff --git a/benchmarking/gat/seastar/egl_gat.py b/benchmarking/gat/seastar/egl_gat.py
index e573c7b..13be530 100644
--- a/benchmarking/gat/seastar/egl_gat.py
+++ b/benchmarking/gat/seastar/egl_gat.py
@@ -1,14 +1,9 @@
 """Torch modules for graph attention networks(GAT)."""
 # pylint: disable= no-member, arguments-differ, invalid-name
-import argparse
-
 import torch
-import dgl
 from torch import nn
-from seastar import Seastar
-from seastar.backend.pytorch_backend import backend_cb
-
-from seastar import Seastar
+from seastar.compiler.backend.pytorch_backend import backend_cb
+from seastar.compiler import Seastar
 
 # pylint: enable=W0235
 class EglGATConv(nn.Module):
@@ -49,7 +44,6 @@ def reset_parameters(self):
             nn.init.xavier_normal_(self.res_fc.weight, gain=gain)
 
     def forward(self, graph, feat):
-        graph = graph.local_var()
         h_dst = h_src = self.feat_drop(feat)
         feat_src = feat_dst = self.fc(h_src).view(-1, self._num_heads, self._out_feats)
         el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
diff --git a/benchmarking/gat/seastar/nn_train.py b/benchmarking/gat/seastar/nn_train.py
new file mode 100644
index 0000000..bd64882
--- /dev/null
+++ b/benchmarking/gat/seastar/nn_train.py
@@ -0,0 +1,277 @@
+"""
+Graph Attention Networks in DGL using SPMV optimization.
+Multiple heads are also batched together for faster training.
+References
+----------
+Paper: https://arxiv.org/abs/1710.10903
+Author's code: https://github.com/PetarV-/GAT
+Pytorch implementation: https://github.com/Diego999/pyGAT
+"""
+
+import argparse
+import time
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from seastar.graph.static.StaticGraph import StaticGraph
+from seastar.dataset.cora import CoraDataset
+from seastar.nn.pytorch.gat_conv import GATConv
+from utils import EarlyStopping
+import snoop
+import numpy as np
+
+
+class EglGAT(nn.Module):
+    def __init__(self,
+                 g,
+                 num_layers,
+                 in_dim,
+                 num_hidden,
+                 num_classes,
+                 heads,
+                 activation,
+                 feat_drop,
+                 attn_drop,
+                 negative_slope,
+                 residual):
+        super(EglGAT, self).__init__()
+        self.g = g
+        self.num_layers = num_layers
+        self.gat_layers = nn.ModuleList()
+        self.activation = activation
+        # input projection (no residual)
+        self.gat_layers.append(GATConv(
+            in_dim, num_hidden, heads[0],
+            feat_drop, attn_drop, negative_slope, False, self.activation))
+        # hidden layers
+        for l in range(1, num_layers):
+            # due to multi-head, the in_dim = num_hidden * num_heads
+            self.gat_layers.append(GATConv(
+                num_hidden * heads[l-1], num_hidden, heads[l],
+                feat_drop, attn_drop, negative_slope, residual, self.activation))
+        # output projection
+        self.gat_layers.append(GATConv(
+            num_hidden * heads[-2], num_classes, heads[-1],
+            feat_drop, attn_drop, negative_slope, residual, None))
+
+    def forward(self, inputs):
+        h = inputs
+        for l in range(self.num_layers):
+            h = self.gat_layers[l](self.g, h).flatten(1)
+        # output projection
+        logits = self.gat_layers[-1](self.g, h).mean(1)
+        return logits
+
+
+def accuracy(logits, labels):
+    _, indices = torch.max(logits, dim=1)
+    correct = torch.sum(indices == labels)
+    return correct.item() * 1.0 / len(labels)
+
+
+def evaluate(model, features, labels, mask):
+    model.eval()
+    with torch.no_grad():
+        logits = model(features)
+        logits = logits[mask]
+        labels = labels[mask]
+        return accuracy(logits, labels)
+
+
+def train(args):
+    # load and preprocess dataset
+    path = '../../dataset/' + 'cora' + '/'
+    '''
+    edges = np.loadtxt(path + 'edges.txt')
+    edges = edges.astype(int)
+
+    features = np.loadtxt(path + 'features.txt')
+
+    train_mask = np.loadtxt(path + 'train_mask.txt')
+    train_mask = train_mask.astype(int)
+
+    labels = np.loadtxt(path + 'labels.txt')
+    labels = labels.astype(int)
+    '''
+    # cora = CoraDataset(verbose=True)
+    # features = torch.FloatTensor(cora.get_all_features())
+    # labels = torch.LongTensor(cora.get_all_targets())
+    # train_mask = cora.get_train_mask()
+    # test_mask = cora.get_test_mask()
+    # train_mask = torch.BoolTensor(train_mask)
+    # test_mask = torch.BoolTensor(test_mask)
+    # num_edges = len(cora.get_edges())
+
+    edges = np.load(path + 'edges.npy')
+    features = np.load(path + 'features.npy')
+    train_mask = np.load(path + 'train_mask.npy')
+    labels = np.load(path + 'labels.npy')
+
+    num_edges = edges.shape[0]
+    num_nodes = features.shape[0]
+    num_feats = features.shape[1]
+    n_classes = int(max(labels) - min(labels) + 1)
+
+    assert train_mask.shape[0] == num_nodes
+
+    print('dataset {}'.format("Cora"))
+    print('# of edges : {}'.format(num_edges))
+    print('# of nodes : {}'.format(num_nodes))
+    print('# of features : {}'.format(num_feats))
+
+    features = torch.FloatTensor(features)
+    labels = torch.LongTensor(labels)
+
+    if hasattr(torch, 'BoolTensor'):
+        train_mask = torch.BoolTensor(train_mask)
+
+    else:
+        train_mask = torch.ByteTensor(train_mask)
+
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        torch.cuda.set_device(args.gpu)
+        features = features.cuda()
+        labels = labels.cuda()
+        train_mask = train_mask.cuda()
+
+    print(edges.shape)
+    edges_lst = [(edge[0],edge[1]) for edge in edges]
+    # print(edges_lst)
+    edge_weight = [1 for _ in range(num_edges)]
+    g = StaticGraph(edges_lst, edge_weight, num_nodes)
+
+    # create model
+    heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
+    model = EglGAT(g,
+                args.num_layers,
+                num_feats,
+                args.num_hidden,
+                n_classes,
+                heads,
+                F.elu,
+                args.in_drop,
+                args.attn_drop,
+                args.negative_slope,
+                args.residual)
+    print(model)
+    if args.early_stop:
+        stopper = EarlyStopping(patience=100)
+    if cuda:
+        model.cuda()
+    loss_fcn = torch.nn.CrossEntropyLoss()
+
+    # use optimizer
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
+
+    # initialize graph
+    dur = []
+    record_time = 0
+    avg_run_time = 0
+    Used_memory = 0
+
+    for epoch in range(args.num_epochs):
+        #print('epoch = ', epoch) 
+        #print('mem0 = {}'.format(mem0))
+        torch.cuda.synchronize()
+        tf = time.time()
+        model.train()
+        if epoch >= 3:
+            t0 = time.time()
+        # forward
+        logits = model(features)
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+        now_mem = torch.cuda.max_memory_allocated(0)
+        print('now_mem : ', now_mem)
+        Used_memory = max(now_mem, Used_memory)
+        tf1 =time.time()
+
+        optimizer.zero_grad()
+        torch.cuda.synchronize()
+        t1 =time.time()
+        loss.backward()
+        optimizer.step()
+        t2 =time.time()
+        run_time_this_epoch = t2 - tf
+
+        if epoch >= 3:
+            dur.append(time.time() - t0)
+            record_time += 1
+
+            avg_run_time += run_time_this_epoch
+
+        train_acc = accuracy(logits[train_mask], labels[train_mask])
+
+        #log for each step
+        print('Epoch {:05d} | Time(s) {:.4f} | train_acc {:.6f} | Used_Memory {:.6f} mb'.format(
+            epoch, run_time_this_epoch, train_acc, (now_mem * 1.0 / (1024**2))
+        ))
+        '''
+        if args.fastmode:
+            val_acc = accuracy(logits[val_mask], labels[val_mask])
+        else:
+            val_acc = evaluate(model, features, labels, val_mask)
+            if args.early_stop:
+                if stopper.step(val_acc, model):   
+                    break
+
+        print("Epoch {:05d} | Time(s) {:.4f} | Loss {:.4f} | TrainAcc {:.4f} |"
+              " ValAcc /{:.4f} | ETputs(KTEPS) {:.2f}".
+              format(epoch, np.mean(dur), loss.item(), train_acc,
+                     val_acc, n_edges / np.mean(dur) / 1000))
+        
+        '''
+    
+
+    if args.early_stop:
+        model.load_state_dict(torch.load('es_checkpoint.pt'))
+
+    #OUTPUT we need
+    avg_run_time = avg_run_time *1. / record_time
+    Used_memory /= (1024**3)
+    print('^^^{:6f}^^^{:6f}'.format(Used_memory, avg_run_time))
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description='GAT')
+    # register_data_args(parser)
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--gpu", type=int, default=0,
+                        help="which GPU to use. Set -1 to use CPU.")
+    parser.add_argument("--num_epochs", type=int, default=200,
+                        help="number of training epochs")
+    parser.add_argument("--num_heads", type=int, default=8,
+                        help="number of hidden attention heads")
+    parser.add_argument("--num_out_heads", type=int, default=1,
+                        help="number of output attention heads")
+    parser.add_argument("--num_layers", type=int, default=1,
+                        help="number of hidden layers")
+    parser.add_argument("--num_hidden", type=int, default=32,
+                        help="number of hidden units")
+    parser.add_argument("--residual", action="store_true", default=False,
+                        help="use residual connection")
+    parser.add_argument("--in_drop", type=float, default=.6,
+                        help="input feature dropout")
+    parser.add_argument("--attn_drop", type=float, default=.6,
+                        help="attention dropout")
+    parser.add_argument("--lr", type=float, default=0.005,
+                        help="learning rate")
+    parser.add_argument('--weight_decay', type=float, default=5e-4,
+                        help="weight decay")
+    parser.add_argument('--negative_slope', type=float, default=0.2,
+                        help="the negative slope of leaky relu")
+    parser.add_argument('--early_stop', action='store_true', default=False,
+                        help="indicates whether to use early stop or not")
+    parser.add_argument('--fastmode', action="store_true", default=False,
+                        help="skip re-evaluate the validation set")
+    args = parser.parse_args()
+
+    print(args)
+        
+    train(args)
diff --git a/benchmarking/gat/seastar/train.py b/benchmarking/gat/seastar/train.py
index 1e7f0cc..fb40a47 100644
--- a/benchmarking/gat/seastar/train.py
+++ b/benchmarking/gat/seastar/train.py
@@ -9,17 +9,15 @@
 """
 
 import argparse
-import numpy as np
-import networkx as nx
 import time
 import torch
 import torch.nn.functional as F
-import dgl
-import dgl.function as fn
+from seastar.graph.static.StaticGraph import StaticGraph
+from seastar.dataset.cora import CoraDataset
 from egl_gat import EglGAT
 from utils import EarlyStopping
-from dgl.data import register_data_args, load_data
 import snoop
+import numpy as np
 
 
 def accuracy(logits, labels):
@@ -39,7 +37,7 @@ def evaluate(model, features, labels, mask):
 
 def train(args):
     # load and preprocess dataset
-    path = '../../dataset/' + str(args.dataset) + '/'
+    path = '../../dataset/' + 'cora' + '/'
     '''
     edges = np.loadtxt(path + 'edges.txt')
     edges = edges.astype(int)
@@ -52,6 +50,15 @@ def train(args):
     labels = np.loadtxt(path + 'labels.txt')
     labels = labels.astype(int)
     '''
+    # cora = CoraDataset(verbose=True)
+    # features = torch.FloatTensor(cora.get_all_features())
+    # labels = torch.LongTensor(cora.get_all_targets())
+    # train_mask = cora.get_train_mask()
+    # test_mask = cora.get_test_mask()
+    # train_mask = torch.BoolTensor(train_mask)
+    # test_mask = torch.BoolTensor(test_mask)
+    # num_edges = len(cora.get_edges())
+
     edges = np.load(path + 'edges.npy')
     features = np.load(path + 'features.npy')
     train_mask = np.load(path + 'train_mask.npy')
@@ -64,7 +71,7 @@ def train(args):
 
     assert train_mask.shape[0] == num_nodes
 
-    print('dataset {}'.format(args.dataset))
+    print('dataset {}'.format("Cora"))
     print('# of edges : {}'.format(num_edges))
     print('# of nodes : {}'.format(num_nodes))
     print('# of features : {}'.format(num_feats))
@@ -87,12 +94,10 @@ def train(args):
         labels = labels.cuda()
         train_mask = train_mask.cuda()
 
-    u = edges[:,0]
-    v = edges[:,1]
-    g = dgl.graph((u,v), num_nodes=num_nodes)
-    # add self loop
-    g = dgl.add_self_loop(g)
-    g = g.to(features.device)
+    print(edges.shape)
+    edges_lst = [(edge[0],edge[1]) for edge in edges]
+    # print(edges_lst)
+    g = StaticGraph(edges_lst, num_nodes)
 
     # create model
     heads = ([args.num_heads] * args.num_layers) + [args.num_out_heads]
@@ -188,7 +193,7 @@ def train(args):
 if __name__ == '__main__':
 
     parser = argparse.ArgumentParser(description='GAT')
-    register_data_args(parser)
+    # register_data_args(parser)
 
     # COMMENT IF SNOOP IS TO BE ENABLED
     snoop.install(enabled=False)
diff --git a/benchmarking/gcn/pyg/gcn.py b/benchmarking/gcn/pyg/gcn.py
new file mode 100644
index 0000000..a6d14ba
--- /dev/null
+++ b/benchmarking/gcn/pyg/gcn.py
@@ -0,0 +1,37 @@
+"""GCN using builtin functions that enables SPMV optimization.
+
+References:
+- Semi-Supervised Classification with Graph Convolutional Networks
+- Paper: https://arxiv.org/abs/1609.02907
+- Code: https://github.com/tkipf/gcn
+"""
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv
+
+
+class PyG_GCN(nn.Module):
+    def __init__(self,
+                 in_feats,
+                 n_hidden,
+                 n_classes,
+                 n_layers):
+        super(PyG_GCN, self).__init__()
+
+        self.layers = nn.ModuleList()
+        # input layer
+        self.layers.append(GCNConv(in_feats, n_hidden, add_self_loops=False))
+        # hidden layers
+        for i in range(n_layers - 1):
+            self.layers.append(GCNConv(n_hidden, n_hidden, add_self_loops=False))
+        # output layer
+        self.layers.append(GCNConv(n_hidden, n_classes, add_self_loops=False))
+
+    def forward(self, features, edge_index):
+        h = features
+        for layer in self.layers:
+            h = layer(h, edge_index)
+            h = F.relu(h)
+        return h
diff --git a/benchmarking/gcn/pyg/train.py b/benchmarking/gcn/pyg/train.py
new file mode 100644
index 0000000..0d472bd
--- /dev/null
+++ b/benchmarking/gcn/pyg/train.py
@@ -0,0 +1,170 @@
+import argparse, time
+import numpy as np
+import torch
+from seastar.dataset.cora import CoraDataset
+
+# from torch_geometric_temporal.nn.recurrent import TGCN
+from gcn import PyG_GCN
+
+import snoop
+
+def accuracy(logits, labels):
+    _, indices = torch.max(logits, dim=1)
+    correct = torch.sum(indices == labels)
+    return correct.item() * 1.0 / len(labels)
+
+def evaluate(model, features, labels, mask):
+    model.eval()
+    with torch.no_grad():
+        logits = model(features)
+        logits = logits[mask]
+        labels = labels[mask]
+        _, indices = torch.max(logits, dim=1)
+        correct = torch.sum(indices == labels)
+        return correct.item() * 1.0 / len(labels)
+
+# GPU | CPU
+def get_default_device():
+    
+    if torch.cuda.is_available():
+        return torch.device('cuda:0')
+    else:
+        return torch.device('cpu')
+
+def to_default_device(data):
+    
+    if isinstance(data,(list,tuple)):
+        return [to_default_device(x,get_default_device()) for x in data]
+    
+    return data.to(get_default_device(),non_blocking = True)
+
+def main(args):
+    
+    cora = CoraDataset(verbose=True)
+    
+    features = torch.FloatTensor(cora.get_all_features())
+    labels = torch.LongTensor(cora.get_all_targets())
+    
+    train_mask = cora.get_train_mask()
+    test_mask = cora.get_test_mask()
+
+    train_mask = torch.BoolTensor(train_mask)
+    test_mask = torch.BoolTensor(test_mask)
+
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        torch.cuda.set_device(args.gpu)
+        features = features.cuda()
+        labels = labels.cuda()
+        train_mask = train_mask.cuda()
+        test_mask = test_mask.cuda()
+
+    # g = StaticGraph(cora.get_edges())
+
+    # normalization
+    # degs = torch.from_numpy(g.in_degrees()).type(torch.int32)
+    # norm = torch.pow(degs, -0.5)
+    # norm[torch.isinf(norm)] = 0
+    
+    # norm = to_default_device(norm)
+    # g.ndata['norm'] = norm.unsqueeze(1)
+
+    num_feats = features.shape[1]
+    n_classes = int(max(labels) - min(labels) + 1)
+    train_edges = to_default_device(torch.from_numpy(np.array(cora.get_edges()).T)).type(torch.int64)
+
+    # print("SHAPE")
+    # print(train_edges.shape)
+    # print(train_edges.dtype)
+
+    # model = EglGCN(g,
+    #             num_feats,
+    #             args.num_hidden,
+    #             n_classes,
+    #             args.num_layers,
+    #             F.relu,
+    #             args.dropout)
+
+    model = PyG_GCN(num_feats,
+                args.num_hidden,
+                n_classes,
+                args.num_layers)
+    
+    if cuda:
+        model.cuda()
+    loss_fcn = torch.nn.CrossEntropyLoss()
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(),
+                                 lr=args.lr,
+                                 weight_decay=args.weight_decay)
+
+    # initialize graph
+    dur = []
+    Used_memory = 0
+
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+        # forward
+        logits = model(features, train_edges)
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        now_mem = torch.cuda.max_memory_allocated(0)
+        Used_memory = max(now_mem, Used_memory)
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        train_acc = accuracy(logits[train_mask], labels[train_mask])
+        print('Epoch {:05d} | Time(s) {:.4f} | train_acc {:.6f} | Used_Memory {:.6f} mb'.format(
+            epoch, run_time_this_epoch, train_acc, (now_mem * 1.0 / (1024**2))
+        ))
+
+    Used_memory /= (1024**3)
+    print('^^^{:6f}^^^{:6f}'.format(Used_memory, np.mean(dur)))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GCN')
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--dropout", type=float, default=0.5,
+            help="dropout probability")
+    parser.add_argument("--dataset", type=str,
+            help="Datset to train your model")
+    parser.add_argument("--gpu", type=int, default=0,
+            help="gpu")
+    parser.add_argument("--lr", type=float, default=1e-2,
+            help="learning rate")
+    parser.add_argument("--num_epochs", type=int, default=200,
+            help="number of training epochs")
+    parser.add_argument("--num_hidden", type=int, default=16,
+            help="number of hidden gcn units")
+    parser.add_argument("--num_layers", type=int, default=1,
+            help="number of hidden gcn layers")
+    parser.add_argument("--weight-decay", type=float, default=5e-4,
+            help="Weight for L2 loss")
+    parser.add_argument("--self-loop", action='store_true',
+            help="graph self-loop (default=False)")
+    parser.set_defaults(self_loop=False)
+    args = parser.parse_args()
+    print(args)
+
+    main(args)
\ No newline at end of file
diff --git a/benchmarking/gcn/seastar/gcn_spmv.py b/benchmarking/gcn/seastar/gcn_spmv.py
index 115d0b1..74bcba2 100644
--- a/benchmarking/gcn/seastar/gcn_spmv.py
+++ b/benchmarking/gcn/seastar/gcn_spmv.py
@@ -8,8 +8,6 @@
 import math
 import torch
 import torch.nn as nn
-import dgl
-import dgl.function as fn
 from seastar.compiler import Seastar
 from seastar.compiler.backend.pytorch.torch_callback import SeastarBackendTorch
 
@@ -38,10 +36,12 @@ def __init__(self,
         self.seastar = Seastar(SeastarBackendTorch())
 
     def reset_parameters(self):
-        stdv = 1. / math.sqrt(self.weight.size(1))
-        self.weight.data.uniform_(-stdv, stdv)
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        # self.weight.data.uniform_(-stdv, stdv)
+        nn.init.xavier_uniform_(self.weight)
         if self.bias is not None:
-            self.bias.data.uniform_(-stdv, stdv)
+            # self.bias.data.uniform_(-stdv, stdv)
+            nn.init.zeros_(self.bias)
 
     def forward(self, h):
         if self.dropout:
diff --git a/benchmarking/gcn/seastar/nn_train.py b/benchmarking/gcn/seastar/nn_train.py
new file mode 100644
index 0000000..9a83f93
--- /dev/null
+++ b/benchmarking/gcn/seastar/nn_train.py
@@ -0,0 +1,238 @@
+import argparse, time
+import numpy as np
+import networkx as nx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from rich import inspect
+from rich.pretty import pprint
+
+from seastar.graph.static.StaticGraph import StaticGraph
+from seastar.dataset.cora import CoraDataset
+
+import snoop
+
+from seastar.nn.pytorch.graph_conv import GraphConv
+import gc
+import sys
+import pynvml
+
+class EglGCN(nn.Module):
+    def __init__(self,
+                 g,
+                 in_feats,
+                 n_hidden,
+                 n_classes,
+                 n_layers,
+                 activation):
+        super(EglGCN, self).__init__()
+        self.g = g
+        self.layers = nn.ModuleList()
+        # input layer
+        self.layers.append(GraphConv(in_feats, n_hidden, activation))
+        # hidden layers
+        for i in range(n_layers - 1):
+            self.layers.append(GraphConv(n_hidden, n_hidden, activation))
+        # output layer
+        self.layers.append(GraphConv(n_hidden, n_classes, None))
+
+    def forward(self, g, features):
+        h = features
+        for layer in self.layers:
+            h = layer(g, h)
+        return h
+
+def accuracy(logits, labels):
+    _, indices = torch.max(logits, dim=1)
+    correct = torch.sum(indices == labels)
+    return correct.item() * 1.0 / len(labels)
+
+def evaluate(model, features, labels, mask):
+    model.eval()
+    with torch.no_grad():
+        logits = model(features)
+        logits = logits[mask]
+        labels = labels[mask]
+        _, indices = torch.max(logits, dim=1)
+        correct = torch.sum(indices == labels)
+        return correct.item() * 1.0 / len(labels)
+
+# GPU | CPU
+def get_default_device():
+    
+    if torch.cuda.is_available():
+        return torch.device('cuda:0')
+    else:
+        return torch.device('cpu')
+
+def to_default_device(data):
+    
+    if isinstance(data,(list,tuple)):
+        return [to_default_device(x,get_default_device()) for x in data]
+    
+    return data.to(get_default_device(),non_blocking = True)
+
+def main(args):
+
+
+    
+    cora = CoraDataset(verbose=True)
+
+    # To account for the initial CUDA Context object for pynvml
+    tmp = StaticGraph([(0,0)], [1], 1)
+    a = torch.tensor([2]).cuda()
+    
+    features = torch.FloatTensor(cora.get_all_features())
+    labels = torch.LongTensor(cora.get_all_targets())
+    
+    train_mask = cora.get_train_mask()
+    test_mask = cora.get_test_mask()
+
+    train_mask = torch.BoolTensor(train_mask)
+    test_mask = torch.BoolTensor(test_mask)
+
+    if args.gpu < 0:
+        cuda = False
+    else:
+        cuda = True
+        torch.cuda.set_device(args.gpu)
+        features = features.cuda()
+        labels = labels.cuda()
+        train_mask = train_mask.cuda()
+        test_mask = test_mask.cuda()
+
+    print("Features Shape: ", features.shape)
+    edge_weight = [1 for _ in range(len(cora.get_edges()))]
+
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
+    g = StaticGraph(cora.get_edges(), edge_weight, features.shape[0])
+    graph_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+
+    # A simple sanity check
+    print("Measuerd Graph Size (pynvml): ", graph_mem, " B")
+    print("Measuerd Graph Size (pynvml): ", (graph_mem)/(1024**2), " MB")
+    
+    # add self loop
+    # if args.self_loop:
+    #     g = dgl.add_self_loop(g)
+    
+    # g = g.to(features.device)
+
+    # normalization
+    degs = torch.from_numpy(g.weighted_in_degrees()).type(torch.int32)
+    norm = torch.pow(degs, -0.5)
+    norm[torch.isinf(norm)] = 0
+    
+    norm = to_default_device(norm)
+    g.set_ndata("norm", norm.unsqueeze(1))
+    # print("Norm Shape: ", g.ndata['norm'].shape)
+
+    num_feats = features.shape[1]
+    n_classes = int(max(labels) - min(labels) + 1)
+    print("Num Classes: ",n_classes)
+
+    model = EglGCN(g,
+                num_feats,
+                args.num_hidden,
+                n_classes,
+                args.num_layers,
+                F.relu)
+    
+    if cuda:
+        model.cuda()
+    loss_fcn = torch.nn.CrossEntropyLoss()
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(),
+                                 lr=args.lr,
+                                 weight_decay=args.weight_decay)
+
+    # initialize graph
+    dur = []
+    Used_memory = 0
+
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        # print(f"\n\n------- BEFORE MODEL E={epoch} ----------")
+        # # gc.collect()
+        # for obj in gc.get_objects():
+        #     if torch.is_tensor(obj):
+        #         print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+
+        # forward
+        logits = model(g, features)
+
+        # print(f"------- AFTER MODEL E={epoch} -------\n\n")
+        # # gc.collect()
+        # for obj in gc.get_objects():
+        #     if torch.is_tensor(obj):
+        #         print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+
+        loss = loss_fcn(logits[train_mask], labels[train_mask])
+        # now_mem = torch.cuda.max_memory_allocated(0)
+        # Used_memory = max(now_mem, Used_memory)
+        # now_mem_pynvml = (
+        #         pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+        #     )
+
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        now_mem = torch.cuda.max_memory_allocated(0) + graph_mem
+        Used_memory = max(now_mem, Used_memory)
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        train_acc = accuracy(logits[train_mask], labels[train_mask])
+        print('Epoch {:05d} | Time(s) {:.4f} | train_acc {:.6f} | Used_Memory {:.6f} mb '.format(
+            epoch, run_time_this_epoch, train_acc, (now_mem * 1.0 / (1024**2))
+        ))
+
+    Used_memory /= (1024**3)
+    print('^^^{:6f}^^^{:6f}'.format(Used_memory, np.mean(dur)))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GCN')
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--dropout", type=float, default=0.5,
+            help="dropout probability")
+    parser.add_argument("--dataset", type=str,
+            help="Datset to train your model")
+    parser.add_argument("--gpu", type=int, default=0,
+            help="gpu")
+    parser.add_argument("--lr", type=float, default=1e-2,
+            help="learning rate")
+    parser.add_argument("--num_epochs", type=int, default=200,
+            help="number of training epochs")
+    parser.add_argument("--num_hidden", type=int, default=16,
+            help="number of hidden gcn units")
+    parser.add_argument("--num_layers", type=int, default=1,
+            help="number of hidden gcn layers")
+    parser.add_argument("--weight-decay", type=float, default=5e-4,
+            help="Weight for L2 loss")
+    parser.add_argument("--self-loop", action='store_true',
+            help="graph self-loop (default=False)")
+    parser.set_defaults(self_loop=False)
+    args = parser.parse_args()
+    print(args)
+
+    main(args)
\ No newline at end of file
diff --git a/benchmarking/gcn/seastar/train.py b/benchmarking/gcn/seastar/train.py
index 95960a8..e7fd3a1 100644
--- a/benchmarking/gcn/seastar/train.py
+++ b/benchmarking/gcn/seastar/train.py
@@ -14,6 +14,9 @@
 import snoop
 
 from gcn_spmv import EglGCN
+import gc
+import sys
+import nvidia_smi
 
 def accuracy(logits, labels):
     _, indices = torch.max(logits, dim=1)
@@ -46,8 +49,15 @@ def to_default_device(data):
     return data.to(get_default_device(),non_blocking = True)
 
 def main(args):
+
+    nvidia_smi.nvmlInit()
+    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
     
     cora = CoraDataset(verbose=True)
+
+    tmp = StaticGraph([(0,0)], [1], 1)
+
+    initial_used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used
     
     features = torch.FloatTensor(cora.get_all_features())
     labels = torch.LongTensor(cora.get_all_targets())
@@ -68,7 +78,9 @@ def main(args):
         train_mask = train_mask.cuda()
         test_mask = test_mask.cuda()
 
-    g = StaticGraph(cora.get_edges())
+    print("Features Shape: ", features.shape)
+    edge_weight = [1 for _ in range(len(cora.get_edges()))]
+    g = StaticGraph(cora.get_edges(), edge_weight, features.shape[0])
     
     # add self loop
     # if args.self_loop:
@@ -77,15 +89,17 @@ def main(args):
     # g = g.to(features.device)
 
     # normalization
-    degs = torch.from_numpy(g.in_degrees()).type(torch.int32)
+    degs = torch.from_numpy(g.weighted_in_degrees()).type(torch.int32)
     norm = torch.pow(degs, -0.5)
     norm[torch.isinf(norm)] = 0
     
     norm = to_default_device(norm)
     g.ndata['norm'] = norm.unsqueeze(1)
+    print("Norm Shape: ", g.ndata['norm'].shape)
 
     num_feats = features.shape[1]
     n_classes = int(max(labels) - min(labels) + 1)
+    print("Num Classes: ",n_classes)
 
     model = EglGCN(g,
                 num_feats,
@@ -113,10 +127,28 @@ def main(args):
         if cuda:
             torch.cuda.synchronize()
         t0 = time.time()
+
+        # print(f"\n\n------- BEFORE MODEL E={epoch} ----------")
+        # # gc.collect()
+        # for obj in gc.get_objects():
+        #     if torch.is_tensor(obj):
+        #         print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+
         # forward
         logits = model(features)
+
+        # print(f"------- AFTER MODEL E={epoch} -------\n\n")
+        # # gc.collect()
+        # for obj in gc.get_objects():
+        #     if torch.is_tensor(obj):
+        #         print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+
         loss = loss_fcn(logits[train_mask], labels[train_mask])
-        now_mem = torch.cuda.max_memory_allocated(0)
+        # now_mem = torch.cuda.max_memory_allocated(0)
+        # Used_memory = max(now_mem, Used_memory)
+        now_mem = (
+                nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+            )
         Used_memory = max(now_mem, Used_memory)
 
         optimizer.zero_grad()
diff --git a/benchmarking/link-pred/pygt/train.py b/benchmarking/link-pred/pygt/train.py
new file mode 100644
index 0000000..0d4eb9d
--- /dev/null
+++ b/benchmarking/link-pred/pygt/train.py
@@ -0,0 +1,155 @@
+import json
+import urllib.request
+import random
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch_geometric_temporal.nn.recurrent import TGCN
+import time
+import argparse
+from seastar.dataset.LinkPredDatasetLoader import LinkPredDatasetLoader
+
+class PyGT_TGCN(torch.nn.Module):
+  def __init__(self, node_features, multiplier):
+    super(PyGT_TGCN, self).__init__()
+    self.temporal = TGCN(node_features, multiplier*node_features)
+    self.linear = torch.nn.Linear(multiplier*node_features, node_features)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    return y, h
+  
+  def decode(self, z, edge_label_index):
+    return (z[edge_label_index[0]] * z[edge_label_index[1]]).sum(
+        dim=-1
+    )  # product of a pair of nodes on each edge
+
+
+def get_default_device():
+    if torch.cuda.is_available():
+        return torch.device('cuda:0')
+    else:
+        return torch.device('cpu')
+
+def to_default_device(data):
+    if isinstance(data,(list,tuple)):
+        return [to_default_device(x,get_default_device()) for x in data]
+    
+    # removed non_blocking
+    # (ORIGINAL) data.to(get_default_device(),non_blocking = True)
+    return data.to(get_default_device())
+
+def main(args):
+
+    dataloader = LinkPredDatasetLoader(args.dataset_dir, args.dataset, args.num_nodes, verbose=True)
+    print("Loaded dataset into the train.py pygt", flush=True)
+    
+    edge_list = dataloader.get_edges()
+    train_edges_lst = edge_list
+    train_features = to_default_device(torch.randn((dataloader._max_num_nodes, args.feat_size)))
+    pos_neg_edges_lst, pos_neg_targets_lst = dataloader.get_pos_neg_edges()
+    print("Features ready", flush=True)
+    
+    model = to_default_device(PyGT_TGCN(args.feat_size, args.multiplier))
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    print("Model ready", flush=True)
+
+    # metrics
+    dur = []
+    cuda = True
+    backprop_every = args.backprop_every
+
+    train_edges_lst = [to_default_device(torch.from_numpy(edge_index)) for edge_index in train_edges_lst]
+    pos_neg_edges_lst = [to_default_device(torch.from_numpy(pos_neg_edges)) for pos_neg_edges in pos_neg_edges_lst]
+    pos_neg_targets_lst = [to_default_device(torch.from_numpy(pos_neg_targets).type(torch.float32)) for pos_neg_targets in pos_neg_targets_lst]
+    print("Edge lists ready", flush=True)
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    # train
+    print("Training...\n")
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+
+        if cuda:
+            torch.cuda.synchronize()
+        
+        t0 = time.time()
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+        gpu_mem_arr = []
+        y_hat = train_features
+
+        # num_iter = int(len(pos_neg_targets) / backprop_every)
+        num_iter = len(pos_neg_targets_lst)
+        for index in range(num_iter): 
+            
+            # for k in range(backprop_every):
+            #     timestamp = index*backprop_every + k
+
+                # if timestamp >= dataloader.total_timestamps:
+                #     continue
+            timestamp = index
+            y_hat, hidden_state = model(y_hat, train_edges_lst[timestamp], None, hidden_state)
+            out = model.decode(y_hat, pos_neg_edges_lst[timestamp]).view(-1)
+            cost = cost + criterion(out, pos_neg_targets_lst[timestamp])
+                
+    
+        cost = cost / (num_iter+1)
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+            
+            # Detaching output from the graph so that it isnt backpropagated
+            # y_hat = y_hat.detach()
+            # hidden_state = hidden_state.detach()
+        
+        if args.measure_space:
+            used_gpu_mem = torch.cuda.max_memory_allocated(0)
+            gpu_mem_arr.append(used_gpu_mem)
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        if args.measure_space:
+            print('Epoch {:03d} | Time(s) {:.4f} | BCE {:.2f} | Used GPU Memory (Max) {:.3f} mb | Used GPU Memory (Avg) {:.3f} mb '.format(
+                epoch, run_time_this_epoch, cost, (max(gpu_mem_arr) * 1.0 / (1024**2)), ((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr)))
+            ))
+        else:
+            print('Epoch {:03d} | Time(s) {:.4f} | BCE {:.2f} '.format(epoch, run_time_this_epoch, cost))
+
+    print('Average Time taken: {:6f}'.format(np.mean(dur)))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GCN')
+
+    parser.add_argument("--lr", type=float, default=1e-2,
+            help="learning rate")
+    parser.add_argument("--feat_size", type=int, default=8,
+            help="feature size")
+    parser.add_argument("--multiplier", type=int, default=2,
+            help="multiplier")
+    parser.add_argument("--num_epochs", type=int, default=1,
+            help="number of training epochs")
+    parser.add_argument("--num_nodes", type=int, default=0,
+            help="number of nodes")
+    parser.add_argument("--backprop_every", type=int, default=1,
+            help="number of timestamps after which backprop should happen")
+    parser.add_argument("--measure_space", action='store_true',
+            help="Measure space")
+    parser.add_argument(
+        "--dataset_dir", type=str, default="wiki-talk", help="dataset directory"
+    )
+    parser.add_argument("--dataset", type=str, default="wiki-talk-temporal",
+            help="Name of the Dataset", metavar="dataset")
+    args = parser.parse_args()
+    print(args)
+
+    main(args)
\ No newline at end of file
diff --git a/benchmarking/link-pred/seastar/train.py b/benchmarking/link-pred/seastar/train.py
new file mode 100644
index 0000000..ea4400a
--- /dev/null
+++ b/benchmarking/link-pred/seastar/train.py
@@ -0,0 +1,213 @@
+import json
+import urllib.request
+import random
+import numpy as np
+import torch
+import torch.nn.functional as F
+import time
+import argparse
+from seastar.dataset.LinkPredDatasetLoader import LinkPredDatasetLoader
+from seastar.nn.pytorch.temporal.tgcn import TGCN
+from seastar.graph.dynamic.gpma.GPMAGraph import GPMAGraph
+from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
+from seastar.graph.dynamic.DynamicGraph import DynamicGraph
+from seastar.graph.dynamic.naive.NaiveGraph import NaiveGraph
+import pynvml
+
+class SeastarTGCN(torch.nn.Module):
+    def __init__(self, node_features, multiplier):
+        super(SeastarTGCN, self).__init__()
+        self.temporal = TGCN(node_features, multiplier*node_features)
+        self.linear = torch.nn.Linear(multiplier*node_features, node_features)
+
+    def forward(self, g, node_feat, edge_weight, hidden_state):
+        h = self.temporal(g, node_feat, edge_weight, hidden_state)
+        y = F.relu(h)
+        y = self.linear(y)
+        return y, h
+  
+    def decode(self, z, edge_label_index):
+        return (z[edge_label_index[0]] * z[edge_label_index[1]]).sum(
+            dim=-1
+        )  # product of a pair of nodes on each edge
+
+
+def get_default_device():
+    if torch.cuda.is_available():
+        return torch.device('cuda:0')
+    else:
+        return torch.device('cpu')
+
+def to_default_device(data):
+    if isinstance(data,(list,tuple)):
+        return [to_default_device(x,get_default_device()) for x in data]
+    
+    # removed non_blocking
+    # (ORIGINAL) data.to(get_default_device(),non_blocking = True)
+    return data.to(get_default_device())
+
+def main(args):
+
+    dataloader = LinkPredDatasetLoader(args.dataset_dir, args.dataset, args.num_nodes, verbose=True, for_seastar=True)
+    print("Loaded dataset into the train.py seastar")
+
+    # to allocate initial CUDA context object
+    Graph = None
+    if args.type == "naive":
+        Graph = NaiveGraph([[(0,0)]],1)
+    elif args.type == "pcsr":
+        Graph = PCSRGraph([[(0,1)]],2) # PCSRGraph([[(0,1)]],2)
+    elif args.type == "gpma":
+        Graph = GPMAGraph([[(0,0)]],1)
+    
+
+    train_edges_lst = dataloader.get_edges()
+    pos_neg_edges_lst, pos_neg_targets_lst = dataloader.get_pos_neg_edges()
+    train_features = to_default_device(torch.randn((dataloader._max_num_nodes, args.feat_size)))
+    edge_weight = None
+    print("Features ready")
+    
+    model = to_default_device(SeastarTGCN(args.feat_size, args.multiplier))
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+    print("Model ready")
+
+    # metrics
+    dur = []
+    cuda = True
+    backprop_every = args.backprop_every
+
+    pos_neg_edges_lst = [to_default_device(torch.from_numpy(pos_neg_edges)) for pos_neg_edges in pos_neg_edges_lst]
+    pos_neg_targets_lst = [to_default_device(torch.from_numpy(pos_neg_targets).type(torch.float32)) for pos_neg_targets in pos_neg_targets_lst]
+    print("Edge lists ready")
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    initial_used_gpu_mem = 0
+    graph_mem = 0
+
+    if args.measure_space:
+        initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
+
+    if args.type == "naive":
+        G = NaiveGraph(train_edges_lst, args.num_nodes)
+    elif args.type == "pcsr":
+        G = PCSRGraph(train_edges_lst, args.num_nodes)
+    elif args.type == "gpma":
+        G = GPMAGraph(train_edges_lst, args.num_nodes)
+    else:
+        print("Error: Invalid Type")
+        quit()
+
+    if args.measure_space:
+        graph_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+        print("Measuerd Graph Size (pynvml): ", graph_mem, " B")
+        print("Measuerd Graph Size (pynvml): ", (graph_mem)/(1024**2), " MB")
+
+    # train
+    print("Training...\n")
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+
+        if cuda:
+            torch.cuda.synchronize()
+        
+        t0 = time.time()
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+        gpu_mem_arr = []
+        y_hat = train_features
+        graph_update_time = 0
+        forward_prop_time = 0
+
+        # num_iter = int(len(pos_neg_targets) / backprop_every)
+        num_iter = len(pos_neg_targets_lst)
+        for index in range(num_iter): 
+            if args.measure_space:
+                initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
+
+            tGstart = time.time()
+            G.get_graph(index)
+            graph_update_time += time.time() - tGstart
+
+            if args.measure_space:
+                graph_mem_delta = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+                graph_mem = graph_mem + graph_mem_delta
+
+            if G.get_ndata("norm") is None:
+                degs = torch.from_numpy(G.in_degrees()).type(torch.float32)
+                norm = torch.pow(degs, -0.5)
+                norm[torch.isinf(norm)] = 0
+                norm = to_default_device(norm)
+                G.set_ndata("norm", norm.unsqueeze(1))
+            
+            # for k in range(backprop_every):
+            #     timestamp = index*backprop_every + k
+
+                # if timestamp >= dataloader.total_timestamps:
+                #     continue
+            timestamp = index
+            y_hat, hidden_state = model(G, y_hat, edge_weight, hidden_state)
+            out = model.decode(y_hat, pos_neg_edges_lst[timestamp]).view(-1)
+            cost = cost + criterion(out, pos_neg_targets_lst[timestamp])
+
+                
+        forward_prop_time = time.time() - t0
+        cost = cost / (num_iter+1)
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+            
+            # Detaching output from the graph so that it isnt backpropagated
+            # y_hat = y_hat.detach()
+            # hidden_state = hidden_state.detach()
+        
+        if args.measure_space:
+            used_gpu_mem = torch.cuda.max_memory_allocated(0) + graph_mem
+            gpu_mem_arr.append(used_gpu_mem)
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        if args.measure_space:
+            print('Epoch {:03d} | Time(s) {:.4f} | BCE {:.2f} | Used GPU Memory (Max) {:.3f} mb | Used GPU Memory (Avg) {:.3f} mb | SeastarGraph (GPU Mem) {:.3f} mb'.format(
+                epoch, run_time_this_epoch, cost, (max(gpu_mem_arr) * 1.0 / (1024**2)), ((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr))), ((graph_mem * 1.0)/(1024**2))
+            ))
+        else:
+            print('Epoch {:03d} | Time(s) {:.4f} | Graph Update Time(s) {:.4f} | Forward Prop Time (s) {:.4f} | BCE {:.2f} '.format(epoch, run_time_this_epoch, graph_update_time, forward_prop_time, cost))
+
+    print('Average Time taken: {:6f}'.format(np.mean(dur)))
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GCN')
+
+    parser.add_argument("--lr", type=float, default=1e-2,
+            help="learning rate")
+    parser.add_argument("--type", type=str, default="naive", help="Seastar Type")
+    parser.add_argument("--feat_size", type=int, default=8,
+            help="feature size")
+    parser.add_argument("--multiplier", type=int, default=2,
+            help="multiplier")
+    parser.add_argument("--num_epochs", type=int, default=1,
+            help="number of training epochs")
+    parser.add_argument("--num_nodes", type=int, default=0,
+            help="number of nodes")
+    parser.add_argument("--backprop_every", type=int, default=1,
+            help="number of timestamps after which backprop should happen")
+    parser.add_argument("--measure_space", action='store_true',
+            help="Measure space")
+    parser.add_argument(
+        "--dataset_dir", type=str, default="wiki-talk", help="dataset directory"
+    )
+    parser.add_argument("--dataset", type=str, default="wiki-talk-temporal",
+            help="Name of the Dataset", metavar="dataset")
+    args = parser.parse_args()
+    print(args)
+
+    main(args)
\ No newline at end of file
diff --git a/benchmarking/tgcn/overall/dynamic_bench.py b/benchmarking/tgcn/overall/dynamic_bench.py
new file mode 100644
index 0000000..9378ec7
--- /dev/null
+++ b/benchmarking/tgcn/overall/dynamic_bench.py
@@ -0,0 +1,479 @@
+import argparse, time
+import numpy as np
+import networkx as nx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from seastar_tgcn import SeastarTGCN
+from pygt_tgcn import PyGT_TGCN
+import snoop
+import os
+
+import nvidia_smi
+import psutil
+
+from rich import inspect
+from rich.pretty import pprint
+from rich.console import Console
+from rich.table import Table
+
+console = Console()
+
+from seastar.graph.dynamic.gpma.GPMAGraph import GPMAGraph
+from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
+from seastar.graph.dynamic.naive.NaiveGraph import NaiveGraph
+from seastar.dataset.FoorahBase import FoorahBase
+
+import seastar.compiler.debugging.print_variables as print_var
+
+# GPU | CPU
+def get_default_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda:0")
+    else:
+        return torch.device("cpu")
+
+
+def to_default_device(data):
+    if isinstance(data, (list, tuple)):
+        return [to_default_device(x, get_default_device()) for x in data]
+
+    return data.to(get_default_device(), non_blocking=True)
+
+
+def run_seastar(dataset_dir, dataset, feat_size, lr, type, max_num_nodes, num_epochs):
+    print_var.is_print_verbose_log = False
+
+    nvidia_smi.nvmlInit()
+    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
+
+    initial_used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used
+    initial_used_cpu_mem = psutil.virtual_memory()[3]
+
+    eng_covid = FoorahBase(dataset_dir, dataset, verbose=False, for_seastar=True)
+
+    edge_list = eng_covid.get_edges()
+    edge_weight_list = eng_covid.get_edge_weights()
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+
+    all_features = to_default_device(torch.FloatTensor(np.array(all_features)))
+    all_targets = to_default_device(torch.FloatTensor(np.array(all_targets)))
+
+    # Hyperparameters
+    train_test_split = 0.8
+
+    # train_test_split for graph
+    train_edges_lst = edge_list[: int(len(edge_list) * train_test_split)]
+    test_edges_lst = edge_list[int(len(edge_list) * train_test_split) :]
+
+    train_edge_weights_lst = edge_weight_list[
+        : int(len(edge_weight_list) * train_test_split)
+    ]
+    test_edge_weights_lst = edge_weight_list[
+        int(len(edge_weight_list) * train_test_split) :
+    ]
+
+    # train_test_split for features
+    train_features = all_features[: int(len(all_features) * train_test_split)]
+    train_targets = all_targets[: int(len(all_targets) * train_test_split)]
+
+    test_features = all_features[int(len(all_features) * train_test_split) :]
+    test_targets = all_targets[int(len(all_targets) * train_test_split) :]
+
+    model = to_default_device(SeastarTGCN(feat_size))
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+
+    # metrics
+    dur = []
+    gpu_move_time_dur = []
+    
+    temp_array = []
+    prop_time_dur = []
+    
+    cuda = True
+
+    if type == "naive":
+        G = NaiveGraph(train_edges_lst, max_num_nodes)
+    elif type == "pcsr":
+        G = PCSRGraph(train_edges_lst, max_num_nodes)
+    elif type == "gpma":
+        G = GPMAGraph(train_edges_lst, max_num_nodes)
+    else:
+        print("Error: Invalid Type")
+        quit()
+
+    # train
+    for epoch in range(num_epochs):
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        G._gpu_move_time = 0
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+
+        gpu_mem_arr = []
+        cpu_mem_arr = []
+
+        update_time_start = G._total_update_time
+
+        # dyn_graph_index is dynamic graph index
+        for index in range(0, len(train_features)):
+            # Getting the graph for a particular timestamp
+            G.get_graph(index)
+
+            # normalization
+            degs = torch.from_numpy(G.in_degrees()).type(torch.int32)
+            norm = torch.pow(degs, -0.5)
+            norm[torch.isinf(norm)] = 0
+            norm = to_default_device(norm)
+            G.ndata["norm"] = norm.unsqueeze(1)
+            edge_weight = to_default_device(
+                torch.FloatTensor(train_edge_weights_lst[index])
+            )
+            edge_weight = torch.unsqueeze(edge_weight, 1)
+
+            # forward propagation
+            y_hat, hidden_state = model(
+                G, train_features[index], edge_weight, hidden_state
+            )
+
+            cost = cost + torch.mean((y_hat - train_targets[index]) ** 2)
+
+            used_gpu_mem = (
+                nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+            )
+            gpu_mem_arr.append(used_gpu_mem)
+            used_cpu_mem = (psutil.virtual_memory()[3]) - initial_used_cpu_mem
+            cpu_mem_arr.append(used_cpu_mem)
+
+        cost = cost / (index + 1)
+
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        run_time_this_epoch = time.time() - t0
+
+        update_time_end = G._total_update_time
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+            gpu_move_time_dur.append(G._gpu_move_time)
+            prop_time_dur.append(update_time_end-update_time_start)  
+                  
+        # temp_array.append(G._total_update_time)
+
+    if G._update_count == 0:
+        time_per_update = 0
+    else:
+        time_per_update = G._total_update_time / G._update_count
+
+    return np.mean(dur), time_per_update, np.mean(gpu_move_time_dur), np.mean(prop_time_dur)
+
+
+def run_pygt(dataset_dir, dataset, feat_size, lr, num_epochs):
+    nvidia_smi.nvmlInit()
+    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
+
+    initial_used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used
+    initial_used_cpu_mem = psutil.virtual_memory()[3]
+
+    eng_covid = FoorahBase(dataset_dir, dataset, verbose=False)
+
+    edge_list = eng_covid.get_edges()
+    edge_weight_list = eng_covid.get_edge_weights()
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+
+    all_features = to_default_device(torch.FloatTensor(np.array(all_features)))
+    all_targets = to_default_device(torch.FloatTensor(np.array(all_targets)))
+
+    # Hyperparameters
+    train_test_split = 0.8
+
+    # train_test_split for graph (Graph)
+    train_edges_lst = edge_list[: int(len(edge_list) * train_test_split)]
+    test_edges_lst = edge_list[int(len(edge_list) * train_test_split) :]
+
+    train_edge_weights_lst = edge_weight_list[
+        : int(len(edge_weight_list) * train_test_split)
+    ]
+    test_edge_weights_lst = edge_weight_list[
+        int(len(edge_weight_list) * train_test_split) :
+    ]
+
+    # train_test_split for features
+    train_features = all_features[: int(len(all_features) * train_test_split)]
+    train_targets = all_targets[: int(len(all_targets) * train_test_split)]
+
+    test_features = all_features[int(len(all_features) * train_test_split) :]
+    test_targets = all_targets[int(len(all_targets) * train_test_split) :]
+
+    model = to_default_device(PyGT_TGCN(feat_size))
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
+
+    # metrics
+    dur = []
+    pygt_gpu_move_dur = []
+    cuda = True
+
+    edge_weight_lst = [
+        to_default_device(torch.FloatTensor(edge_weight))
+        for edge_weight in train_edge_weights_lst
+    ]
+    
+    # NOTE: Previously done like this
+    move_t0 = time.time()
+    train_edges_lst = [
+        to_default_device(torch.from_numpy(np.array(edge_index)))
+        for edge_index in train_edges_lst
+    ]
+    move_t1 = time.time()
+
+    # train
+    for epoch in range(num_epochs):
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+
+        gpu_mem_arr = []
+        cpu_mem_arr = []
+        move_time_total = 0
+
+        # dyn_graph_index is dynamic graph index
+        for index in range(0, len(train_features)):
+            # t1 = time.time()
+
+            edge_weight = edge_weight_lst[index]
+            train_edges = train_edges_lst[index]
+
+            # NOTE: New method for PyG-T
+            # move_t0 = time.time()
+            # train_edges = to_default_device(torch.from_numpy(np.array(train_edges_lst[index])))
+            # move_t1 = time.time()
+            
+            # move_time_total += (move_t1 - move_t0)
+
+            # forward propagation
+            y_hat, hidden_state = model(
+                train_features[index], train_edges, edge_weight, hidden_state
+            )
+            cost = cost + torch.mean((y_hat - train_targets[index]) ** 2)
+
+            used_gpu_mem = (
+                nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+            )
+            gpu_mem_arr.append(used_gpu_mem)
+            used_cpu_mem = (psutil.virtual_memory()[3]) - initial_used_cpu_mem
+            cpu_mem_arr.append(used_cpu_mem)
+
+        cost = cost / (index + 1)
+
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+            # pygt_gpu_move_dur.append(move_time_total)
+
+    time_per_update = 0
+    update_time_epoch = 0
+
+    return np.mean(dur), time_per_update, move_t1-move_t0, 0
+
+
+def main(args):
+    console.log("Starting Benchmarking")
+
+    table = Table(title=f"\nBenchmarking T-GCN\n", show_edge=False, style="black bold")
+
+    table.add_column("Dataset Name", justify="right")
+    table.add_column("Feat. Size", justify="left")
+    table.add_column("Naive", justify="left")
+    table.add_column("PCSR", justify="left")
+    table.add_column("GPMA", justify="left")
+    table.add_column("PyG-T", justify="left")
+
+    update_time_table = Table(
+        title=f"\nTime Taken per Update [black bold](µs)[/black bold]\n", show_edge=False, style="black bold"
+    )
+
+    update_time_table.add_column("Dataset Name", justify="right")
+    update_time_table.add_column("Feat. Size", justify="left")
+    update_time_table.add_column("Naive", justify="left")
+    update_time_table.add_column("PCSR", justify="left")
+    update_time_table.add_column("GPMA", justify="left")
+    update_time_table.add_column("PyG-T", justify="left")
+    
+    gpu_move_time_table = Table(
+        title=f"\nGPU Move Time [black bold](per epoch)[/black bold]\n", show_edge=False, style="black bold"
+    )
+
+    gpu_move_time_table.add_column("Dataset Name", justify="right")
+    gpu_move_time_table.add_column("Feat. Size", justify="left")
+    gpu_move_time_table.add_column("Naive", justify="left")
+    gpu_move_time_table.add_column("PCSR", justify="left")
+    gpu_move_time_table.add_column("GPMA", justify="left")
+    gpu_move_time_table.add_column("PyG-T", justify="left")
+    
+    time_prop_update_table = Table(
+        title=f"\nTotal time taken for update [black bold](per epoch)[/black bold]\n", show_edge=False, style="black bold"
+    )
+
+    time_prop_update_table.add_column("Dataset Name", justify="right")
+    time_prop_update_table.add_column("Feat. Size", justify="left")
+    time_prop_update_table.add_column("Naive", justify="left")
+    time_prop_update_table.add_column("PCSR", justify="left")
+    time_prop_update_table.add_column("GPMA", justify="left")
+    time_prop_update_table.add_column("PyG-T", justify="left")
+
+    dataset_dir = args.dataset_dir
+
+    # creating the dataset list
+    dataset_name = {}
+
+    dataset_feat_size = 8
+    while dataset_feat_size <= args.max_feat_size:
+        dataset_name[f"{dataset_dir}_{dataset_feat_size}"] = {
+            "feat_size": dataset_feat_size,
+            "max_num_nodes": 55000,
+        }
+        dataset_feat_size += 8
+
+    seastar_graph_types = ["naive", "pcsr", "gpma"]
+    # seastar_graph_types = ["pcsr"]
+
+    learning_rate = args.lr
+    num_epochs = args.num_epochs
+
+    for dataset, param in dataset_name.items():
+        results = {}
+        update_time_results = {}
+        gpu_move_time_results = {}
+        total_update_time_results = {}
+
+        # running seastar t-gcn
+        for graph_type in seastar_graph_types:
+            console.log(
+                f"Running [bold yellow]{graph_type}[/bold yellow] on [bold cyan]{dataset}[/bold cyan]"
+            )
+            avg_time, update_time, gpu_move_time, total_update_time = run_seastar(
+                dataset_dir,
+                dataset,
+                param["feat_size"],
+                learning_rate,
+                graph_type,
+                param["max_num_nodes"],
+                num_epochs,
+            )
+            results[graph_type] = round(avg_time, 4)
+            update_time_results[graph_type] = round(update_time * (10**6), 4)
+            gpu_move_time_results[graph_type] = round(gpu_move_time, 4)
+            total_update_time_results[graph_type] = round(total_update_time, 4)
+
+        # running pygt t-gcn
+        console.log(
+            f"Running [bold yellow]PyG-T[/bold yellow] on [bold cyan]{dataset}[/bold cyan]"
+        )
+
+        avg_time, update_time, gpu_move_time, total_update_time = run_pygt(
+            dataset_dir,
+            dataset,
+            param["feat_size"],
+            learning_rate,
+            num_epochs,
+        )
+
+        results["pygt"] = round(avg_time, 4)
+        update_time_results["pygt"] = update_time
+        gpu_move_time_results["pygt"] = round(gpu_move_time, 4)
+        total_update_time_results["pygt"] = round(total_update_time, 4)
+
+        # getting the implementation with the fastest time
+        fast_impl = min(results, key=results.get)
+        results[fast_impl] = f"[bold green]{results[fast_impl]}[/bold green]"
+        
+        fast_gpu_move_time = min(gpu_move_time_results, key=gpu_move_time_results.get)
+        gpu_move_time_results[fast_gpu_move_time] = f"[bold green]{gpu_move_time_results[fast_gpu_move_time]}[/bold green]"
+
+        table.add_row(
+            str(dataset),
+            str(param["feat_size"]),
+            str(results["naive"]),
+            str(results["pcsr"]),
+            str(results["gpma"]),
+            str(results["pygt"]),
+        )
+
+        update_time_table.add_row(
+            str(dataset),
+            str(param["feat_size"]),
+            str(update_time_results["naive"]),
+            str(update_time_results["pcsr"]),
+            str(update_time_results["gpma"]),
+            str(update_time_results["pygt"]),
+        )
+        
+        gpu_move_time_table.add_row(
+            str(dataset),
+            str(param["feat_size"]),
+            str(gpu_move_time_results["naive"]),
+            str(gpu_move_time_results["pcsr"]),
+            str(gpu_move_time_results["gpma"]),
+            str(gpu_move_time_results["pygt"]),
+        )
+        
+        time_prop_update_table.add_row(
+            str(dataset),
+            str(param["feat_size"]),
+            str(total_update_time_results["naive"]),
+            str(total_update_time_results["pcsr"]),
+            str(total_update_time_results["gpma"]),
+            str(total_update_time_results["pygt"]),
+        )
+
+    console.print(table)
+    console.print(update_time_table)
+    console.print(gpu_move_time_table)
+    console.print(time_prop_update_table)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN")
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
+    parser.add_argument(
+        "--dataset_dir", type=str, default="foorah_large", help="dataset directory"
+    )
+    parser.add_argument("--max_feat_size", type=int, default=8, help="max_feature_size")
+    parser.add_argument("--max_num_nodes", type=int, default=1000, help="max_num_nodes")
+    parser.add_argument(
+        "--num_epochs", type=int, default=1, help="number of training epochs"
+    )
+    args = parser.parse_args()
+
+    main(args)
diff --git a/benchmarking/tgcn/overall/pygt_tgcn.py b/benchmarking/tgcn/overall/pygt_tgcn.py
new file mode 100644
index 0000000..b81e145
--- /dev/null
+++ b/benchmarking/tgcn/overall/pygt_tgcn.py
@@ -0,0 +1,16 @@
+import torch
+import torch.nn.functional as F
+from torch_geometric_temporal.nn.recurrent import TGCN
+
+class PyGT_TGCN(torch.nn.Module):
+  def __init__(self, node_features):
+    super(PyGT_TGCN, self).__init__()
+    self.temporal = TGCN(node_features, 2*node_features)
+    self.linear = torch.nn.Linear(2*node_features, node_features)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    
+    return y, h
\ No newline at end of file
diff --git a/benchmarking/tgcn/overall/seastar_tgcn.py b/benchmarking/tgcn/overall/seastar_tgcn.py
new file mode 100644
index 0000000..8c0dff1
--- /dev/null
+++ b/benchmarking/tgcn/overall/seastar_tgcn.py
@@ -0,0 +1,171 @@
+import math
+import torch
+import torch.nn as nn
+from seastar.compiler import Seastar
+import torch.nn.functional as F
+from seastar.compiler.backend.pytorch_backend import backend_cb
+import snoop
+from rich import inspect
+
+from rich.traceback import install
+
+install(show_locals=True)
+
+
+class SeastarGCNLayer(nn.Module):
+    def __init__(self, in_feats, out_feats, activation=None, dropout=None, bias=True):
+        super(SeastarGCNLayer, self).__init__()
+        self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_feats))
+        else:
+            self.bias = None
+        self.activation = activation
+        if dropout:
+            self.dropout = nn.Dropout(p=dropout)
+        else:
+            self.dropout = 0.0
+        self.reset_parameters()
+        self.seastar = Seastar(backend_cb)
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, g, h, edge_weight=None):
+        if self.dropout:
+            h = self.dropout(h)
+
+        h = torch.mm(h, self.weight)
+
+        @self.seastar.compile(nspace=[self, torch])
+        def nb_compute(v):
+            # The nb_edge.src returns a list with one element, this element is an object of NbNode type
+            # hence the translation. Can be cleaned up later.
+            h = sum(
+                [
+                    nb_edge.src[0].norm * nb_edge.src[0].h * nb_edge.weight
+                    for nb_edge in v.inedges
+                ]
+            )
+            h = h * v.norm
+            return h
+
+        h = nb_compute(
+            g=g,
+            n_feats={"norm": g.ndata["norm"], "h": h},
+            e_feats={"weight": edge_weight},
+        )
+
+        # @self.seastar.compile(nspace=[self, torch])
+        # def nb_compute(v):
+        #     # The nb_edge.src returns a list with one element, this element is an object of NbNode type
+        #     # hence the translation. Can be cleaned up later.
+        #     h = sum([nb.h*nb.norm for nb in v.innbs])
+        #     h = h * v.norm
+        #     return h
+        # h = nb_compute(g=g, n_feats={'norm': g.ndata['norm'], 'h' : h})
+
+        # bias
+        if self.bias is not None:
+            h = h + self.bias
+        if self.activation is not None:
+            h = self.activation(h)
+        return h
+
+
+class SeastarTGCNCell(torch.nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+    ):
+        super(SeastarTGCNCell, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # Update GCN Layer
+        self.conv_z = SeastarGCNLayer(self.in_channels, self.out_channels)
+
+        # Update linear layer
+        self.linear_z = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Reset GCN layer
+        self.conv_r = SeastarGCNLayer(self.in_channels, self.out_channels)
+
+        # Reset linear layer
+        self.linear_r = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Candidate (Current Memory Content) GCN layer
+        self.conv_h = SeastarGCNLayer(self.in_channels, self.out_channels)
+
+        # Candidate linear layer
+        self.linear_h = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+    def _set_hidden_state(self, X, H):
+        if H is None:
+            H = torch.zeros(X.shape[0], self.out_channels).to(X.device)
+        return H
+
+    def _calculate_update_gate(self, g, X, edge_weight, H):
+        # print(X.shape)
+        h = self.conv_z(g, X, edge_weight=edge_weight)
+        # print(h.shape)
+        # print(H.shape)
+        # print("\n")
+        # quit()
+        Z = torch.cat((h, H), axis=1)  # axis values need to be checked
+        Z = self.linear_z(Z)
+        Z = torch.sigmoid(Z)
+        return Z
+
+    def _calculate_reset_gate(self, g, X, edge_weight, H):
+        R = torch.cat(
+            (self.conv_r(g, X, edge_weight=edge_weight), H), axis=1
+        )  # axis values need to be checked
+        R = self.linear_r(R)
+        R = torch.sigmoid(R)
+        return R
+
+    def _calculate_candidate_state(self, g, X, edge_weight, H, R):
+        H_tilde = torch.cat(
+            (self.conv_h(g, X, edge_weight=edge_weight), H * R), axis=1
+        )  # axis values need to be checked
+        H_tilde = self.linear_h(H_tilde)
+        H_tilde = torch.tanh(H_tilde)
+        return H_tilde
+
+    def _calculate_hidden_state(self, Z, H, H_tilde):
+        H = Z * H + (1 - Z) * H_tilde
+        return H
+
+    def forward(
+        self,
+        g,
+        X: torch.Tensor,
+        edge_weight: torch.FloatTensor = None,
+        H: torch.FloatTensor = None,
+    ) -> torch.FloatTensor:
+        H = self._set_hidden_state(X, H)
+        Z = self._calculate_update_gate(g, X, edge_weight, H)
+        R = self._calculate_reset_gate(g, X, edge_weight, H)
+        H_tilde = self._calculate_candidate_state(g, X, edge_weight, H, R)
+        H = self._calculate_hidden_state(Z, H, H_tilde)
+        return H
+
+
+class SeastarTGCN(torch.nn.Module):
+    def __init__(self, node_features):
+        super(SeastarTGCN, self).__init__()
+        self.temporal = SeastarTGCNCell(node_features, node_features * 2)
+        self.linear = torch.nn.Linear(node_features * 2, node_features)
+
+    def forward(self, g, node_feat, edge_weight, hidden_state):
+        h = self.temporal(g, node_feat, edge_weight, hidden_state)
+        y = F.relu(h)
+        y = self.linear(y)
+
+        return y, h
diff --git a/benchmarking/tgcn/pygt-dynamic-bench/eng_tgcn.py b/benchmarking/tgcn/pygt-dynamic-bench/eng_tgcn.py
new file mode 100644
index 0000000..91e487f
--- /dev/null
+++ b/benchmarking/tgcn/pygt-dynamic-bench/eng_tgcn.py
@@ -0,0 +1,15 @@
+import torch
+import torch.nn.functional as F
+from torch_geometric_temporal.nn.recurrent import TGCN
+
+class Eng_PyGT_TGCN(torch.nn.Module):
+  def __init__(self, node_features):
+    super(Eng_PyGT_TGCN, self).__init__()
+    self.temporal = TGCN(node_features, 32, add_self_loops=False)
+    self.linear = torch.nn.Linear(32, 1)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    return y, h
\ No newline at end of file
diff --git a/benchmarking/tgcn/pygt-dynamic-bench/tgcn.py b/benchmarking/tgcn/pygt-dynamic-bench/tgcn.py
new file mode 100644
index 0000000..b47d5f4
--- /dev/null
+++ b/benchmarking/tgcn/pygt-dynamic-bench/tgcn.py
@@ -0,0 +1,17 @@
+import torch
+import torch.nn.functional as F
+from torch_geometric_temporal.nn.recurrent import TGCN
+
+class PyGT_TGCN(torch.nn.Module):
+  def __init__(self, node_features):
+    super(PyGT_TGCN, self).__init__()
+    self.temporal = TGCN(node_features, 2*node_features, add_self_loops=False)
+    self.linear = torch.nn.Linear(2*node_features, node_features)
+    # self.temporal = TGCN(node_features, 32)
+    # self.linear = torch.nn.Linear(32, 1)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    return y, h
\ No newline at end of file
diff --git a/benchmarking/tgcn/pygt-dynamic-bench/train.py b/benchmarking/tgcn/pygt-dynamic-bench/train.py
new file mode 100644
index 0000000..29dd7fd
--- /dev/null
+++ b/benchmarking/tgcn/pygt-dynamic-bench/train.py
@@ -0,0 +1,235 @@
+import argparse, time
+import numpy as np
+import networkx as nx
+from seastar.dataset.wikimaths import WikiMaths
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import urllib
+from tqdm import tqdm
+from eng_tgcn import Eng_PyGT_TGCN
+from tgcn import PyGT_TGCN
+import snoop
+import os
+import gc
+
+from rich import inspect
+from rich.pretty import pprint
+
+# from rich.traceback import install
+# install(show_locals=True)
+
+from seastar.dataset.SoorahBase import SoorahBase
+from seastar.dataset.FoorahBase import FoorahBase
+from seastar.dataset.EnglandCOVID import EnglandCOVID
+
+
+# GPU | CPU
+def get_default_device():
+    
+    if torch.cuda.is_available():
+        return torch.device('cuda:0')
+    else:
+        return torch.device('cpu')
+
+def to_default_device(data):
+    
+    if isinstance(data,(list,tuple)):
+        return [to_default_device(x,get_default_device()) for x in data]
+    
+    # removed non_blocking
+    # (ORIGINAL) data.to(get_default_device(),non_blocking = True)
+    return data.to(get_default_device())
+
+def main(args):
+
+    if torch.cuda.is_available():
+        print("🎉 CUDA is available")
+        
+    else:
+        print("😔 CUDA is not available")
+    
+    
+    eng_covid = FoorahBase(args.dataset_dir, args.dataset, verbose=True)
+    # eng_covid = EnglandCOVID(verbose=True)
+    # args.feat_size = 8
+    
+    print("Loaded dataset into the train.py pygt")
+    
+    edge_list = eng_covid.get_edges()
+    edge_weight_list = eng_covid.get_edge_weights()
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+    
+    all_features = to_default_device(torch.FloatTensor(np.array(all_features)))
+    all_targets = to_default_device(torch.FloatTensor(np.array(all_targets)))
+    
+    # print("📝📝📝 GPU Memory in just storing features/targets is: {}".format(((nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem) * 1.0)/(1024**2)))
+    # print("📝📝📝 CPU Memory in just storing features/targets is: {}".format(((psutil.virtual_memory()[3] - initial_used_cpu_mem) * 1.0)/(1024**2)))
+    # print("\n")
+    # Hyperparameters
+    train_test_split = 0.8
+    
+    # train_test_split for graph (Graph)
+    train_edges_lst = edge_list[:int(len(edge_list) * train_test_split)]
+    test_edges_lst = edge_list[int(len(edge_list) * train_test_split):]
+    
+    train_edge_weights_lst = edge_weight_list[:int(len(edge_weight_list) * train_test_split)]
+    test_edge_weights_lst = edge_weight_list[int(len(edge_weight_list) * train_test_split):]
+
+    # train_test_split for features
+    train_features = all_features[:int(len(all_features) * train_test_split)]
+    train_targets = all_targets[:int(len(all_targets) * train_test_split)]
+    
+    test_features = all_features[int(len(all_features) * train_test_split):]
+    test_targets = all_targets[int(len(all_targets) * train_test_split):]
+
+    model = to_default_device(PyGT_TGCN(args.feat_size))
+    # model = to_default_device(Eng_PyGT_TGCN(args.feat_size))
+
+    # print("📝📝📝 GPU Memory after storing model is: {}".format(((nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem) * 1.0)/(1024**2)))
+    # print("📝📝📝 CPU Memory after storing model is: {}".format(((psutil.virtual_memory()[3] - initial_used_cpu_mem) * 1.0)/(1024**2)))
+    # print("\n")
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    # metrics
+    dur = []
+    cuda = True
+    
+    # edge_weight_lst = [to_default_device(torch.FloatTensor(edge_weight)) for edge_weight in train_edge_weights_lst]
+    train_edges_lst = [to_default_device(torch.from_numpy(np.array(edge_index))) for edge_index in train_edges_lst]
+    
+    # print("📝📝📝 GPU Memory after storing edge list/weights: {}".format(((nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem) * 1.0)/(1024**2)))
+    # print("📝📝📝 CPU Memory after storing edge list/weights: {}".format(((psutil.virtual_memory()[3] - initial_used_cpu_mem) * 1.0)/(1024**2)))
+    # print("\n")
+    
+    # train_edges_lst = [to_default_device(torch.from_numpy(np.array(edge_index).T)) for edge_index in train_edges_lst]
+
+    # G = GPMAGraph(train_edges_lst)
+    # G = PCSRGraph(train_edges_lst)
+    # G = NaiveGraph(train_edges_lst)
+
+    # train
+    print("Training...\n")
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+        
+        gpu_mem_arr = []
+        cpu_mem_arr = []
+
+        # dyn_graph_index is dynamic graph index
+        for index in range(0,len(train_features)): 
+            
+            # t1 = time.time()
+            
+            edge_weight = to_default_device(
+                torch.FloatTensor(train_edge_weights_lst[index])
+            )
+            edge_weight = torch.unsqueeze(edge_weight, 1)
+            train_edges = train_edges_lst[index]
+
+            # print(f">>> PRINTING TENSOR TRACE FOR T={index} <<<")
+            # torch.cuda.synchronize()
+            # gc.collect()
+            # for obj in gc.get_objects():
+            #     try:
+            #         if (torch.is_tensor(obj) or ((hasattr(obj, 'data') and torch.is_tensor(obj.data)))) and obj.device != torch.device("cpu"):
+            #             print(type(obj), obj.size(), obj.device)
+            #     except:
+            #         pass
+            # print(f">>> END PRINTING TENSOR TRACE FOR T={index} <<<")
+
+            # forward propagation
+            y_hat, hidden_state = model(train_features[index], train_edges, edge_weight, hidden_state)
+
+            # print(f">>> PRINTING TENSOR TRACE FOR T={index} <<<")
+            # torch.cuda.synchronize()
+            # gc.collect()
+            # for obj in gc.get_objects():
+            #     try:
+            #         if (torch.is_tensor(obj) or ((hasattr(obj, 'data') and torch.is_tensor(obj.data)))) and obj.device != torch.device("cpu"):
+            #             print(type(obj), obj.size(), obj.device)
+            #     except:
+            #         pass
+            # print(f">>> END PRINTING TENSOR TRACE FOR T={index} <<<")
+
+            cost = cost + torch.mean((y_hat-train_targets[index])**2)
+            
+            # run_time_this_timestamp = time.time() - t1
+            # print(f"⌛⌛⌛ Takes a total of {run_time_this_timestamp}")
+
+            # print(f">>> PRINTING TENSOR TRACE FOR T={index} <<<")
+            # torch.cuda.synchronize()
+            # gc.collect()
+            # for obj in gc.get_objects():
+            #     try:
+            #         if (torch.is_tensor(obj) or ((hasattr(obj, 'data') and torch.is_tensor(obj.data)))) and obj.device != torch.device("cpu"):
+            #             print(type(obj), obj.size(), obj.device)
+            #     except:
+            #         pass
+            # print(f">>> END PRINTING TENSOR TRACE FOR T={index} <<<")
+            
+            # if index == 2:
+            #     quit()
+            
+        # print("📝📝📝 GPU Memory after training is: {}".format(((nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem)*1.0))/(1024**2))
+        # print("📝📝📝 GPU Memory after training is: {}".format(((nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem) * 1.0)/(1024**2)))
+        # print("📝📝📝 CPU Memory after training is: {}".format(((psutil.virtual_memory()[3] - initial_used_cpu_mem) * 1.0)/(1024**2)))
+        # print("\n")
+    
+        cost = cost / (index+1)
+        
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+        
+        used_gpu_mem = torch.cuda.max_memory_allocated(0)
+        gpu_mem_arr.append(used_gpu_mem)
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        print('Epoch {:03d} | Time(s) {:.4f} | MSE {:.2f} | Used GPU Memory (Max) {:.3f} mb | Used GPU Memory (Avg) {:.3f} mb '.format(
+            epoch, run_time_this_epoch, cost, (max(gpu_mem_arr) * 1.0 / (1024**2)), ((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr)))
+        ))
+
+    print('Average Time taken: {:6f}'.format(np.mean(dur)))
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='GCN')
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+
+    parser.add_argument("--lr", type=float, default=1e-2,
+            help="learning rate")
+    parser.add_argument("--feat_size", type=int, default=8,
+            help="feature size")
+    parser.add_argument("--num_epochs", type=int, default=1,
+            help="number of training epochs")
+    parser.add_argument(
+        "--dataset_dir", type=str, default="foorah_large", help="dataset directory"
+    )
+    parser.add_argument("--dataset", type=str, default="soorah_base",
+            help="Name of the Soorah Dataset", metavar="dataset")
+    args = parser.parse_args()
+    print(args)
+
+    main(args)
\ No newline at end of file
diff --git a/benchmarking/tgcn/seastar-dynamic-bench/eng_tgcn.py b/benchmarking/tgcn/seastar-dynamic-bench/eng_tgcn.py
new file mode 100644
index 0000000..9dc679c
--- /dev/null
+++ b/benchmarking/tgcn/seastar-dynamic-bench/eng_tgcn.py
@@ -0,0 +1,164 @@
+import math
+import torch
+import torch.nn as nn
+from seastar.compiler import Seastar
+import torch.nn.functional as F
+from seastar.compiler.backend.pytorch_backend import backend_cb
+import snoop
+from rich import inspect
+
+from rich.traceback import install
+install(show_locals=True)
+
+class SeastarGCNLayer(nn.Module):
+    def __init__(self,
+                 in_feats,
+                 out_feats,
+                 activation = None,
+                 dropout = None,
+                 bias=True):
+        super(SeastarGCNLayer, self).__init__()
+        self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_feats))
+        else:
+            self.bias = None
+        self.activation = activation
+        if dropout:
+            self.dropout = nn.Dropout(p=dropout)
+        else:
+            self.dropout = 0.
+        self.reset_parameters()
+        self.seastar = Seastar(backend_cb)
+
+    def reset_parameters(self):
+        # stdv = 1. / math.sqrt(self.weight.size(1))
+        # self.weight.data.uniform_(-stdv, stdv)
+        nn.init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            # self.bias.data.uniform_(-stdv, stdv)
+            nn.init.zeros_(self.bias)
+
+    def forward(self, g, h, edge_weight=None):
+        if self.dropout:
+            h = self.dropout(h)
+
+        h = torch.mm(h, self.weight)
+
+        @self.seastar.compile(nspace=[self, torch])
+        def nb_compute(v):
+            h = sum([nb_edge.src.norm * nb_edge.src.h * nb_edge.edge_weight for nb_edge in v.inedges])
+            h = h * v.norm
+            return h
+        
+        # NOTE: THE VARIABLES THAT ARE PASSED HERE CANT HAVE THE SAME NAME AS ANY OF THE SELF VARIABLES DEFINED IN THIS CLASS
+        # BECAUSE THEN THE VARIABLES IN INPUT_CACHE WOULD GET OVERWRITTEN
+        h = nb_compute(g=g, n_feats={'norm': g.ndata['norm'], 'h' : h}, e_feats={'edge_weight': edge_weight})
+
+        # @self.seastar.compile(nspace=[self, torch])
+        # def nb_compute(v):
+        #     # The nb_edge.src returns a list with one element, this element is an object of NbNode type
+        #     # hence the translation. Can be cleaned up later.
+        #     h = sum([nb.h*nb.norm for nb in v.innbs])
+        #     h = h * v.norm
+        #     return h
+        # h = nb_compute(g=g, n_feats={'norm': g.ndata['norm'], 'h' : h})
+
+        # bias
+        if self.bias is not None:
+            h = h + self.bias
+        if self.activation is not None:
+            h = self.activation(h)
+        return h
+
+class SeastarTGCNCell(torch.nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+    ):
+        super(SeastarTGCNCell, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # Update GCN Layer
+        self.conv_z = SeastarGCNLayer(self.in_channels,self.out_channels)
+        
+        # Update linear layer
+        self.linear_z = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Reset GCN layer
+        self.conv_r = SeastarGCNLayer(self.in_channels,self.out_channels)
+        
+        # Reset linear layer
+        self.linear_r = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Candidate (Current Memory Content) GCN layer
+        self.conv_h = SeastarGCNLayer(self.in_channels,self.out_channels)
+
+        # Candidate linear layer
+        self.linear_h = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+    def _set_hidden_state(self, X, H):
+        if H is None:
+            H = torch.zeros(X.shape[0], self.out_channels).to(X.device)
+        return H
+
+    def _calculate_update_gate(self, g, X, edge_weight, H):
+        # print(X.shape)
+        h = self.conv_z(g, X, edge_weight=edge_weight)
+        # print(h.shape)
+        # print(H.shape)
+        # print("\n")
+        # quit()
+        Z = torch.cat((h, H), axis=1) # axis values need to be checked
+        Z = self.linear_z(Z)
+        Z = torch.sigmoid(Z)
+        return Z
+
+    def _calculate_reset_gate(self, g, X, edge_weight, H):
+        R = torch.cat((self.conv_r(g, X, edge_weight=edge_weight), H), axis=1) # axis values need to be checked
+        R = self.linear_r(R)
+        R = torch.sigmoid(R)
+        return R
+
+    def _calculate_candidate_state(self, g, X, edge_weight, H, R):
+        H_tilde = torch.cat((self.conv_h(g, X, edge_weight=edge_weight), H * R), axis=1) # axis values need to be checked
+        H_tilde = self.linear_h(H_tilde)
+        H_tilde = torch.tanh(H_tilde)
+        return H_tilde
+
+    def _calculate_hidden_state(self, Z, H, H_tilde):
+        H = Z * H + (1 - Z) * H_tilde
+        return H
+
+    def forward(
+        self,
+        g,
+        X: torch.Tensor,
+        edge_weight: torch.FloatTensor = None,
+        H: torch.FloatTensor = None,
+    ) -> torch.FloatTensor:
+
+        H = self._set_hidden_state(X, H)
+        Z = self._calculate_update_gate(g, X, edge_weight, H)
+        R = self._calculate_reset_gate(g, X, edge_weight, H)
+        H_tilde = self._calculate_candidate_state(g, X, edge_weight, H, R)
+        H = self._calculate_hidden_state(Z, H, H_tilde)
+        return H
+        
+
+class Eng_SeastarTGCN(torch.nn.Module):
+  def __init__(self, node_features):
+    super(Eng_SeastarTGCN, self).__init__()
+    self.temporal = SeastarTGCNCell(node_features, 32)
+    self.linear = torch.nn.Linear(32, 1)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    
+    return y, h
\ No newline at end of file
diff --git a/benchmarking/tgcn/seastar-dynamic-bench/nn_train.py b/benchmarking/tgcn/seastar-dynamic-bench/nn_train.py
new file mode 100644
index 0000000..fa689ed
--- /dev/null
+++ b/benchmarking/tgcn/seastar-dynamic-bench/nn_train.py
@@ -0,0 +1,246 @@
+import argparse, time
+import numpy as np
+import torch
+import torch.nn.functional as F
+import snoop
+import pynvml
+from rich.console import Console
+from rich.table import Table
+from seastar.graph.dynamic.gpma.GPMAGraph import GPMAGraph
+from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
+from seastar.graph.dynamic.DynamicGraph import DynamicGraph
+from seastar.graph.dynamic.naive.NaiveGraph import NaiveGraph
+from seastar.dataset.EnglandCOVID import EnglandCOVID
+import seastar.compiler.debugging.print_variables as print_var
+from seastar.nn.pytorch.temporal.tgcn import TGCN
+from seastar.dataset.FoorahBase import FoorahBase
+
+
+class SeastarTGCN(torch.nn.Module):
+    def __init__(self, node_features):
+        super(SeastarTGCN, self).__init__()
+        # self.temporal = TGCN(node_features, 32)
+        # self.linear = torch.nn.Linear(32, 1)
+        self.temporal = TGCN(node_features, 2*node_features)
+        self.linear = torch.nn.Linear(2*node_features, node_features)
+
+    def forward(self, g, node_feat, edge_weight, hidden_state):
+        h = self.temporal(g, node_feat, edge_weight, hidden_state)
+        y = F.relu(h)
+        y = self.linear(y)
+        return y, h
+
+
+# GPU | CPU
+def get_default_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda:0")
+    else:
+        return torch.device("cpu")
+
+
+def to_default_device(data):
+    if isinstance(data, (list, tuple)):
+        return [to_default_device(x, get_default_device()) for x in data]
+
+    # removed non_blocking
+    # (ORIGINAL) data.to(get_default_device(),non_blocking = True)
+    return data.to(get_default_device())
+
+
+def main(args):
+    print_var.is_print_verbose_log = args.verbose
+
+    # preparing the benchmarking table
+    table = Table(title=f"\nBenchmarking T-GCN with {args.type}\n", show_edge=False, style="black bold")
+
+    table.add_column("Epoch", justify="right")
+    table.add_column("Time (s)", justify="left")
+    table.add_column("MSE", justify="left")
+    table.add_column("Max. GPU Memory (mb)", justify="left")
+    table.add_column("Avg. GPU Memory (mb)", justify="left")
+
+
+    if args.type == "gpma":
+        Graph = GPMAGraph([[(0,0)]],1)
+    elif args.type == "naive":
+        Graph = NaiveGraph([[(0,0)]],[[1]],1)
+
+    
+    eng_covid = FoorahBase(args.dataset_dir, args.dataset, verbose=True, for_seastar=True)
+    # eng_covid = EnglandCOVID(for_seastar=True)
+    # args.feat_size = 8
+    # args.max_num_nodes = 129
+
+    print("Loaded dataset into the train.py seastar")
+
+    edge_list = eng_covid.get_edges()
+    edge_weight_list = eng_covid.get_edge_weights()
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+
+    all_features = to_default_device(torch.FloatTensor(np.array(all_features)))
+    all_targets = to_default_device(torch.FloatTensor(np.array(all_targets)))
+
+    # Hyperparameters
+    train_test_split = 0.8
+
+    # train_test_split for graph
+    train_edges_lst = edge_list[: int(len(edge_list) * train_test_split)]
+    test_edges_lst = edge_list[int(len(edge_list) * train_test_split) :]
+
+    train_edge_weights_lst = edge_weight_list[
+        : int(len(edge_weight_list) * train_test_split)
+    ]
+    test_edge_weights_lst = edge_weight_list[
+        int(len(edge_weight_list) * train_test_split) :
+    ]
+
+    # train_test_split for features
+    train_features = all_features[: int(len(all_features) * train_test_split)]
+    train_targets = all_targets[: int(len(all_targets) * train_test_split)]
+
+    test_features = all_features[int(len(all_features) * train_test_split) :]
+    test_targets = all_targets[int(len(all_targets) * train_test_split) :]
+
+    # model = to_default_device(SeastarTGCN(args.feat_size))
+    model = to_default_device(SeastarTGCN(args.feat_size))
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    # metrics
+    dur = []
+    cuda = True
+
+    pynvml.nvmlInit()
+    handle = pynvml.nvmlDeviceGetHandleByIndex(0)
+    initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
+    if args.type == "naive":
+        G = NaiveGraph(train_edges_lst, train_edge_weights_lst, args.max_num_nodes)
+    elif args.type == "pcsr":
+        G = PCSRGraph(train_edges_lst,args.max_num_nodes)
+    elif args.type == "gpma":
+        G = GPMAGraph(train_edges_lst,args.max_num_nodes)
+    else:
+        print("Error: Invalid Type")
+        quit()
+    
+    graph_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+    print("Measuerd Graph Size (pynvml): ", graph_mem, " B")
+    print("Measuerd Graph Size (pynvml): ", (graph_mem)/(1024**2), " MB")
+
+    # A simple sanity check
+    print("Measuerd Graph Size (pynvml): ", graph_mem, " B")
+    print("Measuerd Graph Size (pynvml): ", (graph_mem)/(1024**2), " MB")
+
+    # train
+    print("Training...\n")
+    for epoch in range(args.num_epochs):
+        torch.cuda.reset_peak_memory_stats(0)
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+
+        gpu_mem_arr = []
+
+        # dyn_graph_index is dynamic graph index
+        for index in range(0, len(train_features)):
+
+            if isinstance(G, DynamicGraph):
+                initial_used_gpu_mem = pynvml.nvmlDeviceGetMemoryInfo(handle).used
+
+            # Getting the graph for a particular timestamp
+            G.get_graph(index)
+
+            if isinstance(G, DynamicGraph):
+                graph_mem_delta = pynvml.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+                graph_mem = graph_mem + graph_mem_delta
+                print(f"(E={epoch}|T={index}) Measuerd Graph Size (pynvml): {graph_mem} B")
+                print(f"E={epoch}|T={index}) Measuerd Graph Size (pynvml):  {(graph_mem/(1024**2))} MB")
+
+            # normalization
+            # degs = torch.from_numpy(G.weighted_in_degrees()).type(torch.float32)
+            degs = torch.from_numpy(G.in_degrees()).type(torch.float32)
+            norm = torch.pow(degs, -0.5)
+            norm[torch.isinf(norm)] = 0
+            norm = to_default_device(norm)
+            G.ndata["norm"] = norm.unsqueeze(1)
+            edge_weight = to_default_device(
+                torch.FloatTensor(train_edge_weights_lst[index])
+            )
+            edge_weight = torch.unsqueeze(edge_weight, 1)
+
+            # forward propagation
+            y_hat, hidden_state = model(G, train_features[index], edge_weight, hidden_state)
+
+            cost = cost + torch.mean((y_hat - train_targets[index]) ** 2)
+
+        cost = cost / (index + 1)
+        cost.backward()
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        used_gpu_mem = torch.cuda.max_memory_allocated(0) + graph_mem
+        gpu_mem_arr.append(used_gpu_mem)
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        table.add_row(
+            str(epoch),
+            str(round(run_time_this_epoch, 4)),
+            str(round(cost.item(), 4)),
+            str(round((max(gpu_mem_arr) * 1.0 / (1024**2)), 4)),
+            str(round(((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr))), 4)),
+        )
+
+    console = Console()
+    console.print(table)
+
+    print("\nAverage Time taken (s): {:4f}".format(np.mean(dur)))
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN")
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
+    parser.add_argument("--feat_size", type=int, default=8, help="feature_size")
+    parser.add_argument("--max_num_nodes", type=int, default=1000, help="max_num_nodes")
+    parser.add_argument("--type", type=str, default="naive", help="Seastar Type")
+    parser.add_argument(
+        "--num_epochs", type=int, default=1, help="number of training epochs"
+    )
+    
+    parser.add_argument(
+        "--dataset_dir", type=str, default="foorah_large", help="dataset directory"
+    )
+    
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="foorah_large_8",
+        help="Name of the Soorah Dataset",
+        metavar="dataset",
+    )
+    parser.add_argument(
+        "--verbose",
+        type=bool,
+        default=False,
+        help="If set to true, will print out logs while Seastar compiles your model",
+        metavar="verbose",
+    )
+    args = parser.parse_args()
+
+    main(args)
diff --git a/benchmarking/tgcn/seastar-dynamic-bench/tgcn.py b/benchmarking/tgcn/seastar-dynamic-bench/tgcn.py
new file mode 100644
index 0000000..3850294
--- /dev/null
+++ b/benchmarking/tgcn/seastar-dynamic-bench/tgcn.py
@@ -0,0 +1,164 @@
+import math
+import torch
+import torch.nn as nn
+from seastar.compiler import Seastar
+import torch.nn.functional as F
+from seastar.compiler.backend.pytorch_backend import backend_cb
+import snoop
+from rich import inspect
+
+from rich.traceback import install
+install(show_locals=True)
+
+class SeastarGCNLayer(nn.Module):
+    def __init__(self,
+                 in_feats,
+                 out_feats,
+                 activation = None,
+                 dropout = None,
+                 bias=True):
+        super(SeastarGCNLayer, self).__init__()
+        self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_feats))
+        else:
+            self.bias = None
+        self.activation = activation
+        if dropout:
+            self.dropout = nn.Dropout(p=dropout)
+        else:
+            self.dropout = 0.
+        self.reset_parameters()
+        self.seastar = Seastar(backend_cb)
+
+    def reset_parameters(self):
+        stdv = 1. / math.sqrt(self.weight.size(1))
+        self.weight.data.uniform_(-stdv, stdv)
+        if self.bias is not None:
+            self.bias.data.uniform_(-stdv, stdv)
+
+    def forward(self, g, h, edge_weight=None):
+        if self.dropout:
+            h = self.dropout(h)
+
+        h = torch.mm(h, self.weight)
+
+        @self.seastar.compile(nspace=[self, torch])
+        def nb_compute(v):
+            h = sum([nb_edge.src.norm * nb_edge.src.h * nb_edge.edge_weight for nb_edge in v.inedges])
+            h = h * v.norm
+            return h
+        
+        # NOTE: THE VARIABLES THAT ARE PASSED HERE CANT HAVE THE SAME NAME AS ANY OF THE SELF VARIABLES DEFINED IN THIS CLASS
+        # BECAUSE THEN THE VARIABLES IN INPUT_CACHE WOULD GET OVERWRITTEN
+        h = nb_compute(g=g, n_feats={'norm': g.ndata['norm'], 'h' : h}, e_feats={'edge_weight': edge_weight})
+
+        # @self.seastar.compile(nspace=[self, torch])
+        # def nb_compute(v):
+        #     # The nb_edge.src returns a list with one element, this element is an object of NbNode type
+        #     # hence the translation. Can be cleaned up later.
+        #     h = sum([nb.h*nb.norm for nb in v.innbs])
+        #     h = h * v.norm
+        #     return h
+        # h = nb_compute(g=g, n_feats={'norm': g.ndata['norm'], 'h' : h})
+
+        # bias
+        if self.bias is not None:
+            h = h + self.bias
+        if self.activation is not None:
+            h = self.activation(h)
+        return h
+
+class SeastarTGCNCell(torch.nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+    ):
+        super(SeastarTGCNCell, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # Update GCN Layer
+        self.conv_z = SeastarGCNLayer(self.in_channels,self.out_channels)
+        
+        # Update linear layer
+        self.linear_z = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Reset GCN layer
+        self.conv_r = SeastarGCNLayer(self.in_channels,self.out_channels)
+        
+        # Reset linear layer
+        self.linear_r = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Candidate (Current Memory Content) GCN layer
+        self.conv_h = SeastarGCNLayer(self.in_channels,self.out_channels)
+
+        # Candidate linear layer
+        self.linear_h = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+    def _set_hidden_state(self, X, H):
+        if H is None:
+            H = torch.zeros(X.shape[0], self.out_channels).to(X.device)
+        return H
+
+    def _calculate_update_gate(self, g, X, edge_weight, H):
+        # print(X.shape)
+        h = self.conv_z(g, X, edge_weight=edge_weight)
+        # print(h.shape)
+        # print(H.shape)
+        # print("\n")
+        # quit()
+        Z = torch.cat((h, H), axis=1) # axis values need to be checked
+        Z = self.linear_z(Z)
+        Z = torch.sigmoid(Z)
+        return Z
+
+    def _calculate_reset_gate(self, g, X, edge_weight, H):
+        R = torch.cat((self.conv_r(g, X, edge_weight=edge_weight), H), axis=1) # axis values need to be checked
+        R = self.linear_r(R)
+        R = torch.sigmoid(R)
+        return R
+
+    def _calculate_candidate_state(self, g, X, edge_weight, H, R):
+        H_tilde = torch.cat((self.conv_h(g, X, edge_weight=edge_weight), H * R), axis=1) # axis values need to be checked
+        H_tilde = self.linear_h(H_tilde)
+        H_tilde = torch.tanh(H_tilde)
+        return H_tilde
+
+    def _calculate_hidden_state(self, Z, H, H_tilde):
+        H = Z * H + (1 - Z) * H_tilde
+        return H
+
+    def forward(
+        self,
+        g,
+        X: torch.Tensor,
+        edge_weight: torch.FloatTensor = None,
+        H: torch.FloatTensor = None,
+    ) -> torch.FloatTensor:
+
+        H = self._set_hidden_state(X, H)
+        Z = self._calculate_update_gate(g, X, edge_weight, H)
+        R = self._calculate_reset_gate(g, X, edge_weight, H)
+        H_tilde = self._calculate_candidate_state(g, X, edge_weight, H, R)
+        H = self._calculate_hidden_state(Z, H, H_tilde)
+        return H
+        
+
+class SeastarTGCN(torch.nn.Module):
+  def __init__(self, node_features):
+    super(SeastarTGCN, self).__init__()
+    self.temporal = SeastarTGCNCell(node_features, node_features*2)
+    self.linear = torch.nn.Linear(node_features*2, node_features)
+    # self.temporal = SeastarTGCNCell(node_features, 32)
+    # self.linear = torch.nn.Linear(32, 1)
+
+  def forward(self, g, node_feat, edge_weight, hidden_state):
+    h = self.temporal(g, node_feat, edge_weight, hidden_state)
+    y = F.relu(h)
+    y = self.linear(y)
+    
+    return y, h
\ No newline at end of file
diff --git a/benchmarking/tgcn/seastar-dynamic-bench/train.py b/benchmarking/tgcn/seastar-dynamic-bench/train.py
new file mode 100644
index 0000000..ce2868c
--- /dev/null
+++ b/benchmarking/tgcn/seastar-dynamic-bench/train.py
@@ -0,0 +1,356 @@
+import argparse, time
+import numpy as np
+import networkx as nx
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import urllib
+from tqdm import tqdm
+from tgcn import SeastarTGCN
+from eng_tgcn import Eng_SeastarTGCN
+import snoop
+import os
+import gc
+
+import nvidia_smi
+import psutil
+
+from rich import inspect
+from rich.pretty import pprint
+from rich.console import Console
+from rich.table import Table
+import sys
+
+# from rich.traceback import install
+# install(show_locals=True)
+
+from seastar.graph.dynamic.gpma.GPMAGraph import GPMAGraph
+from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
+from seastar.graph.dynamic.naive.NaiveGraph import NaiveGraph
+from seastar.dataset.EnglandCOVID import EnglandCOVID
+from seastar.dataset.SoorahBase import SoorahBase
+from seastar.dataset.FoorahBase import FoorahBase
+from seastar.dataset.wikimaths import WikiMaths
+
+import seastar.compiler.debugging.print_variables as print_var
+
+
+# GPU | CPU
+def get_default_device():
+    if torch.cuda.is_available():
+        return torch.device("cuda:0")
+    else:
+        return torch.device("cpu")
+
+
+def to_default_device(data):
+    if isinstance(data, (list, tuple)):
+        return [to_default_device(x, get_default_device()) for x in data]
+
+    # removed non_blocking
+    # (ORIGINAL) data.to(get_default_device(),non_blocking = True)
+    return data.to(get_default_device())
+
+
+def main(args):
+    print_var.is_print_verbose_log = args.verbose
+
+    # preparing the benchmarking table
+    table = Table(title=f"\nBenchmarking T-GCN with {args.type}\n", show_edge=False, style="black bold")
+
+    table.add_column("Epoch", justify="right")
+    table.add_column("Time (s)", justify="left")
+    table.add_column("MSE", justify="left")
+    table.add_column("Max. GPU Memory (mb)", justify="left")
+    table.add_column("Avg. GPU Memory (mb)", justify="left")
+    table.add_column("Max. CPU Memory (mb)", justify="left")
+    table.add_column("Avg. CPU Memory (mb)", justify="left")
+
+    nvidia_smi.nvmlInit()
+    handle = nvidia_smi.nvmlDeviceGetHandleByIndex(0)
+
+    if args.type == "gpma":
+        Graph = GPMAGraph([[(0,0)]],1)
+    elif args.type == "naive":
+        Graph = NaiveGraph([[(0,0)]],[[1]],1)
+
+    torch.cuda.empty_cache()
+    initial_used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used
+    initial_used_cpu_mem = psutil.virtual_memory()[3]
+
+    
+    # eng_covid = FoorahBase(args.dataset_dir, args.dataset, verbose=True, for_seastar=True)
+    eng_covid = EnglandCOVID(for_seastar=True)
+    args.feat_size = 8
+    args.max_num_nodes = 129
+
+    print("Loaded dataset into the train.py seastar")
+
+    edge_list = eng_covid.get_edges()
+    edge_weight_list = eng_covid.get_edge_weights()
+    all_features = eng_covid.get_all_features()
+    all_targets = eng_covid.get_all_targets()
+
+    all_features = to_default_device(torch.FloatTensor(np.array(all_features)))
+    all_targets = to_default_device(torch.FloatTensor(np.array(all_targets)))
+
+    torch.cuda.empty_cache()
+    used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+    print(f"STORAGE USED AFTER STORING THE FEATURES: {(used_gpu_mem * 1.0) / (1024**2)}\n")
+
+    # Hyperparameters
+    train_test_split = 0.8
+
+    # train_test_split for graph
+    train_edges_lst = edge_list[: int(len(edge_list) * train_test_split)]
+    test_edges_lst = edge_list[int(len(edge_list) * train_test_split) :]
+
+    train_edge_weights_lst = edge_weight_list[
+        : int(len(edge_weight_list) * train_test_split)
+    ]
+    test_edge_weights_lst = edge_weight_list[
+        int(len(edge_weight_list) * train_test_split) :
+    ]
+
+    # train_test_split for features
+    train_features = all_features[: int(len(all_features) * train_test_split)]
+    train_targets = all_targets[: int(len(all_targets) * train_test_split)]
+
+    test_features = all_features[int(len(all_features) * train_test_split) :]
+    test_targets = all_targets[int(len(all_targets) * train_test_split) :]
+
+    # model = to_default_device(SeastarTGCN(args.feat_size))
+    model = to_default_device(Eng_SeastarTGCN(args.feat_size))
+
+    torch.cuda.empty_cache()
+    used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+    print(f"STORAGE USED AFTER STORING THE MODEL: {(used_gpu_mem * 1.0) / (1024**2)}\n")
+
+    # use optimizer
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
+
+    # metrics
+    dur = []
+    cuda = True
+
+    if args.type == "naive":
+        G = NaiveGraph(train_edges_lst, train_edge_weights_lst, args.max_num_nodes)
+    elif args.type == "pcsr":
+        G = PCSRGraph(train_edges_lst,args.max_num_nodes)
+    elif args.type == "gpma":
+        G = GPMAGraph(train_edges_lst,args.max_num_nodes)
+    else:
+        print("Error: Invalid Type")
+        quit()
+    
+    torch.cuda.empty_cache()
+    used_gpu_mem = nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+    print(f"STORAGE USED AFTER INITIALIZING SEASTAR GRAPH: {(used_gpu_mem * 1.0) / (1024**2)}\n")
+
+    # inspect(G.graph_updates)
+    # num_edges = 0
+    # for key, val in G.graph_updates.items():
+    #     num_edges += len(val["add"])
+    #     num_edges -= len(val["delete"])
+    #     print("For timestamp={} NUM_ADDITIONS={} NUM_DELETIONS={} | TOTAL_EDGE_COUNT={}".format(key,len(val["add"]),len(val["delete"]),num_edges),flush=True)
+    # quit()
+
+    # train
+    print("Training...\n")
+    for epoch in range(args.num_epochs):
+        model.train()
+        if cuda:
+            torch.cuda.synchronize()
+        t0 = time.time()
+
+        cost = 0
+        hidden_state = None
+        optimizer.zero_grad()
+
+        gpu_mem_arr = []
+        cpu_mem_arr = []
+
+        # dyn_graph_index is dynamic graph index
+        for index in range(0, len(train_features)):
+            # t1 = time.time()
+
+            # Getting the graph for a particular timestamp
+            G.get_graph(index)
+
+            # normalization
+            degs = torch.from_numpy(G.weighted_in_degrees()).type(torch.float32)
+            norm = torch.pow(degs, -0.5)
+            norm[torch.isinf(norm)] = 0
+            norm = to_default_device(norm)
+            G.ndata["norm"] = norm.unsqueeze(1)
+            edge_weight = to_default_device(
+                torch.FloatTensor(train_edge_weights_lst[index])
+            )
+            edge_weight = torch.unsqueeze(edge_weight, 1)
+
+            # print(f">>> (Before Model Call) PRINTING TENSOR TRACE FOR T={index} <<<")
+            # torch.cuda.synchronize()
+            # gc.collect()
+            # for obj in gc.get_objects():
+            #     try:
+            #         if (torch.is_tensor(obj) or ((hasattr(obj, 'data') and torch.is_tensor(obj.data)))) and obj.device != torch.device("cpu"):
+            #             print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+            #     except:
+            #         pass
+            # print(f">>> (Before Model Call) END PRINTING TENSOR TRACE FOR T={index} <<<")
+
+            # forward propagation
+            y_hat, hidden_state = model(
+                G, train_features[index], edge_weight, hidden_state
+            )
+
+            # print(f">>> (After Model Call) PRINTING TENSOR TRACE FOR T={index} <<<")
+            # torch.cuda.synchronize()
+            # gc.collect()
+            # for obj in gc.get_objects():
+            #     try:
+            #         if (torch.is_tensor(obj) or ((hasattr(obj, 'data') and torch.is_tensor(obj.data)))) and obj.device != torch.device("cpu"):
+            #             print(type(obj), obj.size(), obj.device, sys.getrefcount(obj))
+            #     except:
+            #         pass
+            # print(f">>> (After Model call) END PRINTING TENSOR TRACE FOR T={index} <<<")
+
+            cost = cost + torch.mean((y_hat - train_targets[index]) ** 2)
+
+            torch.cuda.empty_cache()
+            used_gpu_mem = (
+                nvidia_smi.nvmlDeviceGetMemoryInfo(handle).used - initial_used_gpu_mem
+            )
+            # print(f"(E={epoch} , T={index}) GPU Mem: {(used_gpu_mem * 1.0) / (1024**2)}\n")
+            gpu_mem_arr.append(used_gpu_mem)
+            used_cpu_mem = (psutil.virtual_memory()[3]) - initial_used_cpu_mem
+            cpu_mem_arr.append(used_cpu_mem)
+            # run_time_this_timestamp = time.time() - t1
+            # print(f"⌛⌛⌛ Takes a total of {run_time_this_timestamp}")
+
+            # if index == 2:
+            #     quit()
+
+        cost = cost / (index + 1)
+
+        # t1 = time.time()
+        # print("⚠️⚠️⚠️ Starting Backprop")
+        cost.backward()
+        # print("⚠️⚠️⚠️ Backprop Completed")
+        # print(f"⌛⌛⌛ Time taken for backprop {time.time() - t1}")
+
+        # if epoch == 1:
+        #     quit()
+
+        optimizer.step()
+
+        if cuda:
+            torch.cuda.synchronize()
+
+        run_time_this_epoch = time.time() - t0
+
+        if epoch >= 3:
+            dur.append(run_time_this_epoch)
+
+        table.add_row(
+            str(epoch),
+            str(round(run_time_this_epoch, 4)),
+            str(round(cost.item(), 4)),
+            str(round((max(gpu_mem_arr) * 1.0 / (1024**2)), 4)),
+            str(round(((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr))), 4)),
+            str(round((max(cpu_mem_arr) * 1.0 / (1024**2)), 4)),
+            str(round(((sum(cpu_mem_arr) * 1.0) / ((1024**2) * len(cpu_mem_arr))), 4)),
+        )
+
+        # print(
+        #     "Epoch {:03d} | Time(s) {:.4f} | MSE {:.2f} | Used GPU Memory (Max) {:.3f} mb | Used GPU Memory (Avg) {:.3f} mb | Used CPU Memory (Max) {:.3f} mb | Used CPU Memory (Avg) {:.3f} mb".format(
+        #         epoch,
+        #         run_time_this_epoch,
+        #         cost,
+        #         (max(gpu_mem_arr) * 1.0 / (1024**2)),
+        #         ((sum(gpu_mem_arr) * 1.0) / ((1024**2) * len(gpu_mem_arr))),
+        #         (max(cpu_mem_arr) * 1.0 / (1024**2)),
+        #         ((sum(cpu_mem_arr) * 1.0) / ((1024**2) * len(cpu_mem_arr))),
+        #     )
+        # )
+
+    console = Console()
+    console.print(table)
+
+    print("\nAverage Time taken (s): {:4f}".format(np.mean(dur)))
+
+    # updates_per_sec = round((G._update_count / G._total_update_time), 0)
+    # print(f'Updates per second: {updates_per_sec}')
+
+    # evaluate
+    # print("Evaluating...\n")
+    # model.eval()
+    # cost = 0
+
+    # test_graph_log_dict, test_max_num_nodes = preprocess_graph_structure(test_edges_lst)
+    # G = GPMAGraph(test_graph_log_dict,test_max_num_nodes)
+
+    # # G = PCSRGraph(test_graph_log_dict,test_max_num_nodes)
+
+    # predictions = []
+    # true_y = []
+    # hidden_state=None
+    # # dyn_graph_index is dynamic graph index
+    # for index in range(len(test_features)):
+    #     # normalization
+    #     # degs = G.in_degrees().float()
+    #     degs = torch.from_numpy(G.in_degrees())
+    #     norm = torch.pow(degs, -0.5)
+    #     norm[torch.isinf(norm)] = 0
+    #     norm = to_default_device(norm)
+    #     G.ndata['norm'] = norm.unsqueeze(1)
+    #     edge_weight = to_default_device(torch.FloatTensor(test_edge_weights_lst[index]))
+    #     edge_weight = torch.unsqueeze(edge_weight,1)
+
+    #     # forward propagation
+    #     y_hat, hidden_state = model(G, test_features[index], edge_weight, hidden_state)
+    #     cost = cost + torch.mean((y_hat-test_targets[index])**2)
+    #     predictions.append(y_hat)
+    #     true_y.append(test_targets[index])
+
+    # cost = cost / (index+1)
+    # cost = cost.item()
+    # print("MSE: {:.4f}".format(cost))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="GCN")
+
+    # COMMENT IF SNOOP IS TO BE ENABLED
+    snoop.install(enabled=False)
+
+    parser.add_argument("--lr", type=float, default=1e-2, help="learning rate")
+    parser.add_argument("--feat_size", type=int, default=8, help="feature_size")
+    parser.add_argument("--max_num_nodes", type=int, default=1000, help="max_num_nodes")
+    parser.add_argument("--type", type=str, default="naive", help="Seastar Type")
+    parser.add_argument(
+        "--num_epochs", type=int, default=1, help="number of training epochs"
+    )
+    
+    parser.add_argument(
+        "--dataset_dir", type=str, default="foorah_large", help="dataset directory"
+    )
+    
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        default="foorah_large_8",
+        help="Name of the Soorah Dataset",
+        metavar="dataset",
+    )
+    parser.add_argument(
+        "--verbose",
+        type=bool,
+        default=False,
+        help="If set to true, will print out logs while Seastar compiles your model",
+        metavar="verbose",
+    )
+    args = parser.parse_args()
+
+    main(args)
diff --git a/requirements.txt b/requirements.txt
index 3363467..97fd5a1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,10 +2,8 @@ cuda_python==12.1.0
 Jinja2==3.1.2
 networkx==3.1
 numpy==1.23.4
-nvidia_ml_py3==7.352.0
 pandas==1.3.5
 prettytable==3.7.0
-psutil==5.9.4
 pynvrtc==9.2
 rich==13.3.5
 snoop==0.4.2
@@ -14,4 +12,5 @@ sympy==1.11.1
 termcolor==2.3.0
 tqdm==4.64.1
 pybind11==2.10.4
-build==0.10.0 
\ No newline at end of file
+build==0.10.0 
+pynvml==11.5.0
\ No newline at end of file
diff --git a/seastar/compiler/backend/kernel_wrapper.py b/seastar/compiler/backend/kernel_wrapper.py
index e808eb2..f91b301 100644
--- a/seastar/compiler/backend/kernel_wrapper.py
+++ b/seastar/compiler/backend/kernel_wrapper.py
@@ -3,11 +3,16 @@ def __init__(self):
         pass
     
     @staticmethod
-    def forward(ctx, executor, kid, kernel_args, rets, *args):
-        ctx.backward_cache = executor, kid
+    def forward(executor, kid, kernel_args, rets, *args):
         ret = executor.forward_cb(kid, kernel_args, rets, args)
         return ret
 
+    @staticmethod
+    def setup_context(ctx, inputs, output):
+        # Saving executor and kid in backward cache
+        ctx.backward_cache = inputs[0], inputs[1]
+        ctx.set_materialize_grads(False)
+
     @staticmethod
     def backward(ctx, *gradout):
         executor, kid = ctx.backward_cache
diff --git a/seastar/compiler/code_gen/cuda_check.py b/seastar/compiler/code_gen/cuda_check.py
index 899c3b5..612e8c0 100644
--- a/seastar/compiler/code_gen/cuda_check.py
+++ b/seastar/compiler/code_gen/cuda_check.py
@@ -1,12 +1,17 @@
 from cuda import cuda
-from prettytable import PrettyTable, HEADER, NONE, SINGLE_BORDER
+from prettytable import PrettyTable
 from termcolor import colored
 
 from seastar.compiler.code_gen.cuda_driver import *
 from ctypes import *
 
+NONE = 2
+HEADER = 3
+SINGLE_BORDER = 16
+
 global_gpu_device = None
 
+
 class GPUDevice:
     def __init__(self, print_log=False):
         # breakpoint()
@@ -26,7 +31,7 @@ def __init__(self, print_log=False):
         cuInit(0)
 
         # Getting the number of compatible GPU devices
-        err, self.nGpus = cuDeviceGetCount()    
+        err, self.nGpus = cuDeviceGetCount()
 
         # Going to get certain parameters for only the first GPU
         # i.e ordinal = 0 for cuDeviceGet
@@ -40,16 +45,19 @@ def __init__(self, print_log=False):
         err, self.cc_major = cuDeviceGetAttribute(COMPUTE_CAPABILITY_MAJOR, self.device)
         err, self.cc_minor = cuDeviceGetAttribute(COMPUTE_CAPABILITY_MINOR, self.device)
         err, self.cores = cuDeviceGetAttribute(MULTIPROCESSOR_COUNT, self.device)
-        err, self.thread_per_core = cuDeviceGetAttribute(MAX_THREAD_PER_MULTIPROCESSOR, self.device)
+        err, self.thread_per_core = cuDeviceGetAttribute(
+            MAX_THREAD_PER_MULTIPROCESSOR, self.device
+        )
         err, self.gpu_clockrate = cuDeviceGetAttribute(CLOCK_RATE, self.device)
-        err, self.memory_clockrate = cuDeviceGetAttribute(MEMORY_CLOCK_RATE, self.device)
+        err, self.memory_clockrate = cuDeviceGetAttribute(
+            MEMORY_CLOCK_RATE, self.device
+        )
         err, self.freeMem, self.totalMem = cuMemGetInfo()
 
         if print_log:
             self.log()
 
     def log(self):
-
         log_table = PrettyTable()
         log_table.field_names = ["Device Property", "Value"]
         log_table.align["Device Property"] = "r"
@@ -62,18 +70,43 @@ def log(self):
 
         log_table.add_row(["Number of Devices", self.nGpus])
         log_table.add_row(["Name", self.name])
-        log_table.add_row(["Compute Capability", str(self.cc_major) + "." + str(self.cc_minor)])
+        log_table.add_row(
+            ["Compute Capability", str(self.cc_major) + "." + str(self.cc_minor)]
+        )
         log_table.add_row(["Multiprocessor Count", self.cores])
-        log_table.add_row(["Concurrent Threads", self.cores*self.thread_per_core])
-        log_table.add_row(["GPU Clock", str(self.gpu_clockrate/1000) + colored(" MHz", "dark_grey")])
-        log_table.add_row(["Memory Clock", str(self.memory_clockrate/1000) + colored(" MHz", "dark_grey")])
-        log_table.add_row(["Total Memory", str(self.totalMem/1024**2) + colored(" MiB", "dark_grey")])
-        log_table.add_row(["Free Memory", str(self.freeMem/1024**2) + colored(" MiB", "dark_grey")])
+        log_table.add_row(["Concurrent Threads", self.cores * self.thread_per_core])
+        log_table.add_row(
+            ["GPU Clock", str(self.gpu_clockrate / 1000) + colored(" MHz", "dark_grey")]
+        )
+        log_table.add_row(
+            [
+                "Memory Clock",
+                str(self.memory_clockrate / 1000) + colored(" MHz", "dark_grey"),
+            ]
+        )
+        log_table.add_row(
+            [
+                "Total Memory",
+                str(self.totalMem / 1024**2) + colored(" MiB", "dark_grey"),
+            ]
+        )
+        log_table.add_row(
+            [
+                "Free Memory",
+                str(self.freeMem / 1024**2) + colored(" MiB", "dark_grey"),
+            ]
+        )
 
         print("\n")
         print(log_table)
-        print(colored("\nNote: In case either Total Memory or Free Memory is showing 0\n      it is because no context has been loaded into device", "dark_grey"))
+        print(
+            colored(
+                "\nNote: In case either Total Memory or Free Memory is showing 0\n      it is because no context has been loaded into device",
+                "dark_grey",
+            )
+        )
         print("\n")
 
+
 if __name__ == "__main__":
-    device = GPUDevice(print_log=True)
\ No newline at end of file
+    device = GPUDevice(print_log=True)
diff --git a/seastar/compiler/code_gen/device_info.py b/seastar/compiler/code_gen/device_info.py
index 8c65005..74e825a 100644
--- a/seastar/compiler/code_gen/device_info.py
+++ b/seastar/compiler/code_gen/device_info.py
@@ -1,14 +1,19 @@
 import os
 
 from cuda import cuda
-from prettytable import PrettyTable, HEADER, NONE, SINGLE_BORDER
+from prettytable import PrettyTable
 from termcolor import colored
 
 from seastar.compiler.code_gen.cuda_driver import *
 from ctypes import *
 
+NONE = 2
+HEADER = 3
+SINGLE_BORDER = 16
+
 global_gpu_device = None
 
+
 class DeviceInfo:
     def __init__(self, print_log=False):
         # breakpoint()
@@ -24,16 +29,16 @@ def __init__(self, print_log=False):
         self.totalMem = 0
 
         # Retrieving the path of nvcc
-        self.nvcc_path = os.popen('which nvcc').read()[:-1]
+        self.nvcc_path = os.popen("which nvcc").read()[:-1]
         if not len(self.nvcc_path):
-            self.nvcc_path = '/usr/local/cuda/bin/nvcc'
+            self.nvcc_path = "/usr/local/cuda/bin/nvcc"
 
         # Make sure to call cuInit(), otherwise we won't be
         # able to make any CUDA Driver API calls
         # cuInit(0)
 
         # Getting the number of compatible GPU devices
-        err, self.nGpus = cuDeviceGetCount()    
+        err, self.nGpus = cuDeviceGetCount()
 
         # Going to get certain parameters for only the first GPU
         # i.e ordinal = 0 for cuDeviceGet
@@ -47,16 +52,19 @@ def __init__(self, print_log=False):
         err, self.cc_major = cuDeviceGetAttribute(COMPUTE_CAPABILITY_MAJOR, self.device)
         err, self.cc_minor = cuDeviceGetAttribute(COMPUTE_CAPABILITY_MINOR, self.device)
         err, self.cores = cuDeviceGetAttribute(MULTIPROCESSOR_COUNT, self.device)
-        err, self.thread_per_core = cuDeviceGetAttribute(MAX_THREAD_PER_MULTIPROCESSOR, self.device)
+        err, self.thread_per_core = cuDeviceGetAttribute(
+            MAX_THREAD_PER_MULTIPROCESSOR, self.device
+        )
         err, self.gpu_clockrate = cuDeviceGetAttribute(CLOCK_RATE, self.device)
-        err, self.memory_clockrate = cuDeviceGetAttribute(MEMORY_CLOCK_RATE, self.device)
+        err, self.memory_clockrate = cuDeviceGetAttribute(
+            MEMORY_CLOCK_RATE, self.device
+        )
         err, self.freeMem, self.totalMem = cuMemGetInfo()
 
         if print_log:
             self.log()
 
     def log(self):
-
         log_table = PrettyTable()
         log_table.field_names = ["Device Property", "Value"]
         log_table.align["Device Property"] = "r"
@@ -69,23 +77,48 @@ def log(self):
 
         log_table.add_row(["Number of Devices", self.nGpus])
         log_table.add_row(["Name", self.name])
-        log_table.add_row(["Compute Capability", str(self.cc_major) + "." + str(self.cc_minor)])
+        log_table.add_row(
+            ["Compute Capability", str(self.cc_major) + "." + str(self.cc_minor)]
+        )
         log_table.add_row(["Multiprocessor Count", self.cores])
-        log_table.add_row(["Concurrent Threads", self.cores*self.thread_per_core])
-        log_table.add_row(["GPU Clock", str(self.gpu_clockrate/1000) + colored(" MHz", "dark_grey")])
-        log_table.add_row(["Memory Clock", str(self.memory_clockrate/1000) + colored(" MHz", "dark_grey")])
-        log_table.add_row(["Total Memory", str(self.totalMem/1024**2) + colored(" MiB", "dark_grey")])
-        log_table.add_row(["Free Memory", str(self.freeMem/1024**2) + colored(" MiB", "dark_grey")])
+        log_table.add_row(["Concurrent Threads", self.cores * self.thread_per_core])
+        log_table.add_row(
+            ["GPU Clock", str(self.gpu_clockrate / 1000) + colored(" MHz", "dark_grey")]
+        )
+        log_table.add_row(
+            [
+                "Memory Clock",
+                str(self.memory_clockrate / 1000) + colored(" MHz", "dark_grey"),
+            ]
+        )
+        log_table.add_row(
+            [
+                "Total Memory",
+                str(self.totalMem / 1024**2) + colored(" MiB", "dark_grey"),
+            ]
+        )
+        log_table.add_row(
+            [
+                "Free Memory",
+                str(self.freeMem / 1024**2) + colored(" MiB", "dark_grey"),
+            ]
+        )
 
         print("\n")
         print(log_table)
-        print(colored("\nNote: In case either Total Memory or Free Memory is showing 0\n      it is because no context has been loaded into device", "dark_grey"))
+        print(
+            colored(
+                "\nNote: In case either Total Memory or Free Memory is showing 0\n      it is because no context has been loaded into device",
+                "dark_grey",
+            )
+        )
         print("\n")
 
+
 # def get_global_gpu_device():
 #     if global_gpu_device == None:
 #         return DeviceInfo(print_log=False)
 #     return global_gpu_device
 
 if __name__ == "__main__":
-    device = DeviceInfo(print_log=True)
\ No newline at end of file
+    device = DeviceInfo(print_log=True)
diff --git a/seastar/compiler/code_gen/templates/fa/tpl_fa_gpma.jinja b/seastar/compiler/code_gen/templates/fa/tpl_fa_gpma.jinja
index 6a9c7bc..34fe5d3 100644
--- a/seastar/compiler/code_gen/templates/fa/tpl_fa_gpma.jinja
+++ b/seastar/compiler/code_gen/templates/fa/tpl_fa_gpma.jinja
@@ -25,17 +25,20 @@ extern "C" __global__ void {{kernel_name}}(
 
             for ({{index_type}} e=beg;e<end;++e) {
 
-                unsigned long long {{col_index}} = __ldg(column_indices + e);
+                unsigned long long col_indices_key = __ldg(column_indices + e);
 
                 // GPMA indexes edges starting from 1
                 // Seastar requires edgs to be indexed from 0
                 unsigned int eid = __ldg(eids + e) - 1;
                 
-                unsigned long long mask = (unsigned long long){{row_offset}} << 32;
-                unsigned int dst_check = ({{col_index}} - mask);
-                {{col_index}} = ({{col_index}} - mask);
+                // UNIMPLEMENTED CHECK: Note if the value of col_index exceeds that 
+                // of an int then during casting of unsigned int to int there 
+                // will be errors
+                unsigned int {{col_index}} = (col_indices_key & 0xffffffff);
 
-                if(dst_check != 0xFFFFFFFF && eid != 0){
+                // KEY_MAX: 0xFFFFFFFFFFFFFFFE
+                // COL_IDX_NONE: 0xFFFFFFFF
+                if(col_indices_key != 0xFFFFFFFFFFFFFFFE && {{col_index}} != 0xFFFFFFFF && eid != 0){
                     {{init_inner_offset}}
 
                     {%for edge_stmt in edges%}
diff --git a/seastar/compiler/code_gen/templates/fa/tpl_fa_pcsr.jinja b/seastar/compiler/code_gen/templates/fa/tpl_fa_pcsr.jinja
index 9142b0c..0d6610f 100644
--- a/seastar/compiler/code_gen/templates/fa/tpl_fa_pcsr.jinja
+++ b/seastar/compiler/code_gen/templates/fa/tpl_fa_pcsr.jinja
@@ -28,11 +28,8 @@ extern "C" __global__ void {{kernel_name}}
                 {{index_type}} {{col_index}} = __ldg(column_indices + e);
 
                 // PCSR indexes edges starting from 1
-                // Seastar requires edgs to be indexed from 0
-                // {{index_type}} eid = __ldg(eids + e) - 1;
-
-                // New PCSR
-                {{index_type}} eid = __ldg(eids + e);
+                // Seastar requires edges to be indexed from 0
+                {{index_type}} eid = __ldg(eids + e) - 1;
                 
                 {{init_inner_offset}}
                 
diff --git a/seastar/compiler/execution_unit.py b/seastar/compiler/execution_unit.py
index a451325..d570fdb 100644
--- a/seastar/compiler/execution_unit.py
+++ b/seastar/compiler/execution_unit.py
@@ -270,8 +270,6 @@ def reset_graph_info(self, graph):
             col_indices_ptr = graph.fwd_column_indices_ptr
             eids_ptr = graph.fwd_eids_ptr
         else:
-            #TODO: Will probably have to change this so that this accesses 
-            #backward row_offset, col_indices, eids
             row_offsets_ptr = graph.bwd_row_offset_ptr
             col_indices_ptr = graph.bwd_column_indices_ptr
             eids_ptr = graph.bwd_eids_ptr
@@ -368,10 +366,6 @@ def run(self, tensor_list):
             ASSERT_DRV(ret)
         except Exception as e:
             raise e
-        
-        # sync_ret = cudaDeviceSynchronize()
-        # if ret:
-        #     raise Exception('cudaDeviceSynchronize', sync_ret)
 
 class V2Kernel(Kernel):
     def __init__(self, num_nodes, row_offsets_ptr, col_indices_ptr, eids_ptr, max_dims, kernel_name, compiled_module, launch_config, tile_sizes):
diff --git a/seastar/compiler/executor.py b/seastar/compiler/executor.py
index 2b5e70f..345c2d4 100644
--- a/seastar/compiler/executor.py
+++ b/seastar/compiler/executor.py
@@ -3,6 +3,7 @@
 from collections import deque
 from ..graph.dynamic.DynamicGraph import DynamicGraph
 from seastar.compiler.debugging.seastar_logger import print_log
+import torch
 
 class Stack:
     def __init__(self, val=None):
@@ -36,10 +37,15 @@ def __init__(self):
         # contains arg tensors for the current execution of nb_compute only
         self.current_tensor_map = {}
 
+        # contains IDs of args required for backward propagation of the current timestamp
+        self.bwd_common_tensor_list = []
+        self.bunit_arg_ids = []
+
         self.dep_map = {}
         self.executed_bunit = set()
     
     def reset(self, input_map, f_merged_units, bunits):
+        # print("ENTERING RESET")
         self.dep_map = {}
         #for mu in f_merged_units:
         #    if mu.compiled():
@@ -53,6 +59,18 @@ def reset(self, input_map, f_merged_units, bunits):
                     self.dep_map[arg.id] += 1
                 else:
                     self.dep_map[arg.id] = 1
+
+        # Initializing bwd_common_tensor_list
+        # print("BACKWARD KERNEL ARGS: ", [bu.kernel_args() for bu in bunits])
+        self.bwd_common_tensor_list = []
+        self.bunit_arg_ids = []
+        for bu in bunits:
+            for arg in bu.kernel_args():
+                self.bunit_arg_ids.append(arg.id)
+                if arg.id in input_map.keys():
+                    self.bwd_common_tensor_list.append(arg.id)
+        # print("End of initializing bwd_common_tensor_list\n")
+
         #print('dependency map', self.dep_map)
         self.num_bunits = len(bunits)
         # deletes all tensors that were previously stored here (verified)
@@ -69,8 +87,12 @@ def all_bu_executed(self):
         return len(self.executed_bunit) == self.num_bunits
     
     def track_tensor(self, key, val):
-            self.current_tensor_map[key] = val
+        self.current_tensor_map[key] = val
+        if key in self.bunit_arg_ids:
+            self.bwd_common_tensor_list.append(key)
     
+    def clear_current_tensor_state(self):
+        self.current_tensor_map = {}
     # def clear_cache(self):
     #     rmv_list = []
     #     for k in self.tensor_map:
@@ -204,6 +226,7 @@ def merge_units(self, exec_units):
         return grouped_unit
   
     def restart(self, input_map, graph=None):
+        # print("ENTERING RESTART")
         self.ts.reset(input_map, self.forward_exec_units, self.bulist)
         if graph != None:
             
@@ -265,15 +288,16 @@ def create_tensor_for_vars(self, var_list):
         ret_tensors = {var.id : self.new_zeros(size=[self.num_edges if var.is_edgevar() else self.num_nodes] + list(var.var_shape),
                                                dtype=var.var_dtype,
                                                device=var.device,
-                                               requires_grad=var.requires_grad) for var in var_list if var.id not in self.ts.current_tensor_map}
-        self.ts.current_tensor_map = {**self.ts.current_tensor_map, **ret_tensors}
+                                               requires_grad=False) for var in var_list if var.id not in self.ts.current_tensor_map}
+        
+        for key, val in ret_tensors.items():
+            self.ts.track_tensor(key,val)
 
-    @snoop
     def create_tensor_for_grad_vars(self, var_list, tensor_map):
         ret_tensors = {var.id : self.new_zeros(size=[self.num_edges if var.is_edgevar() else self.num_nodes] + list(var.var_shape),
                                                dtype=var.var_dtype,
                                                device=var.device,
-                                               requires_grad=var.requires_grad) for var in var_list if var.id not in tensor_map}
+                                               requires_grad=False) for var in var_list if var.id not in tensor_map}
         tensor_map = {**tensor_map, **ret_tensors}
         return tensor_map
 
@@ -295,23 +319,27 @@ def execute_compiled(self, uid, FuncWrapper):
         for i,ret in enumerate(rets):
             self.ts.track_tensor(ret.id, ret_tensors[i])
     
-    @snoop
     def forward_cb(self, uid, kernel_args, rets, tensor_list):
         '''FuncWrapper will call this function in forward pass'''
         units = self.forward_exec_units[uid]
         for i,unit in enumerate(units):
             self.execute_unit(unit, [tensor_list[tidx] for tidx in kernel_args[i]])
-            
-        self.ts.tensor_map_stack.push(self.ts.current_tensor_map)
+        
+        self.ts.tensor_map_stack.push({key: self.ts.current_tensor_map[key] for key in self.ts.bwd_common_tensor_list})
+        # self.ts.tensor_map_stack.push(self.ts.current_tensor_map)
         
         if isinstance(self.graph,DynamicGraph):
             self.ts.graph_timestamp_stack.push(self.graph.current_timestamp)
+        
+        ret = tuple([self.ts.current_tensor_map[ret.id] for ret in rets])
+        self.ts.clear_current_tensor_state()
+        return ret
 
-        return tuple([self.ts.current_tensor_map[ret.id] for ret in rets])
-
-    @snoop
     def backward_cb(self, kid, grad_list):
         '''FuncWrapper will call this function in backward pass'''
+
+        # print("BACKWARD CALLED")
+
         # which backward kernel to call? un-executed kernel that has all dependency satisfied. 
         # We need to get the grad_map in order to properly set the variables in compiled kernels.
         funits = self.forward_exec_units[kid]
@@ -358,6 +386,8 @@ def backward_cb(self, kid, grad_list):
         return ret
 
     def execute_prog(self, units):
+        current_tensor_map = self.ts.current_tensor_map
+        self.ts.clear_current_tensor_state()
         for unit in  units:
             for stmt in unit.program:
-                self.ts.track_tensor(stmt.ret.id, stmt.execute([self.ts.current_tensor_map[arg.id] if not is_const_scalar(arg) else arg for arg in stmt.args]))
\ No newline at end of file
+                self.ts.track_tensor(stmt.ret.id, stmt.execute([current_tensor_map[arg.id] if not is_const_scalar(arg) else arg for arg in stmt.args]))
\ No newline at end of file
diff --git a/seastar/compiler/node.py b/seastar/compiler/node.py
index 636ae29..3291332 100644
--- a/seastar/compiler/node.py
+++ b/seastar/compiler/node.py
@@ -22,8 +22,8 @@ class CentralNode(object):
     def __init__(self):
         self.innbs = [NbNode(self, EdgeDirection.IN)] 
         self.outnbs = [NbNode(self, EdgeDirection.OUT)] 
-        self.inedges = [NbEdge(self, EdgeDirection.IN, self.innbs)] 
-        self.outedges = [NbEdge(self, EdgeDirection.OUT, self.outnbs)] 
+        self.inedges = [NbEdge(self, EdgeDirection.IN, self.innbs[i]) for i in range(len(self.innbs))] 
+        self.outedges = [NbEdge(self, EdgeDirection.OUT, self.outnbs[i]) for i in range(len(self.outnbs))] 
     
     def update_allnode(self, feat_map):
         for k,v in feat_map.items():
diff --git a/seastar/compiler/seastar.py b/seastar/compiler/seastar.py
index 9334032..e928d66 100644
--- a/seastar/compiler/seastar.py
+++ b/seastar/compiler/seastar.py
@@ -11,6 +11,8 @@
 from .code_gen import code_gen 
 from .executor import Executor
 from .utils import var_prefix, cen_attr_postfix, inb_attr_postfix
+import gc
+import torch
 
 from seastar.compiler.backend.callback import SeastarBackend
 from seastar.compiler.utils import ValType
@@ -53,13 +55,12 @@ def _setup_executor(self, **kwargs):
             # print('TracedProgram' + str(fprog), 'Ret value:', ret)
             # pretty_print_GIR(fprog,"TGCN GIR")
             self._executor_cache = self._diff_then_compile(ret, fprog, graph)
+        
         for k, v in node_feats.items():
             self._input_cache[var_prefix + k + cen_attr_postfix] = v
             self._input_cache[var_prefix + k + inb_attr_postfix] = v
         for k, v in edge_feats.items():
             self._input_cache[var_prefix+k] = v
-        # print("🔴 Input Cache")
-        # print(self._input_cache)
         self._executor_cache.restart(self._input_cache, graph)
         self._entry_count += 1
         return self._executor_cache
@@ -72,6 +73,8 @@ def _trace(self, nfeats, efeats, input_cache, fprog):
         self._monkey_patch_namespace(old_libs, input_cache, fprog, backend)
         ret = self._f(central_node)
         self._remove_patch(old_libs, backend)
+        self._destroy_central_node(central_node, nfeats, efeats)
+
         if ret == None:
             raise NameError('Ret is none. Execution is aborted')
         return [ret.var] if not isinstance(ret, Iterable) else ret.var
@@ -108,6 +111,17 @@ def _init_central_node(self, nfeats, efeats, fprog, backend):
                     edge_val = self.val_factory.create(ValType.EDGE, v, backend, id=k, fprog=fprog, reduce_dim=True)
                     setattr(e, k, edge_val)
         return cen
+
+    def _destroy_central_node(self, cen, nfeats, efeats):
+        if nfeats:
+            for k, _ in nfeats.items():
+                delattr(cen, k)
+                for n in cen.innbs:
+                    delattr(n, k)
+        if efeats:
+            for k, _ in efeats.items():
+                for e in cen.inedges:
+                    delattr(e, k)
     
     def _monkey_patch_namespace(self, old_libs, input_cache, fprog, backend):
         """Symbolizing central node and its innbs and inedges""" 
diff --git a/seastar/dataset/EnglandCOVID.py b/seastar/dataset/EnglandCOVID.py
index 6ca8013..3239e7b 100644
--- a/seastar/dataset/EnglandCOVID.py
+++ b/seastar/dataset/EnglandCOVID.py
@@ -16,7 +16,7 @@
 install(show_locals=True)
 
 class EnglandCOVID:
-    def __init__(self, verbose: bool = False, lags: int = 8, split=0.75) -> None:
+    def __init__(self, verbose: bool = False, lags: int = 8, split=0.75, for_seastar=False) -> None:
         self.name = "EnglandCOVID"
         self.lags = lags
         self.split = split
@@ -31,9 +31,13 @@ def __init__(self, verbose: bool = False, lags: int = 8, split=0.75) -> None:
 
         self._load_dataset()
         self.total_timestamps = self._dataset["time_periods"]
-        self._get_edge_info()
         self._get_targets_and_features()
-        self._presort_edge_weights()
+
+        if for_seastar:
+            self._get_edge_info_seastar()
+            self._presort_edge_weights()
+        else:
+            self._get_edge_info_pygt()
 
     def get_graph_data(self):
         return self._graph_updates, self._max_num_nodes
@@ -47,13 +51,13 @@ def _load_dataset(self) -> None:
         # THIS NEEDS TO BE EDITED
         # with open('../../dataset/eng_covid/eng_covid.json', 'w') as f:
         #     json.dump(self._dataset,f)
-
-    def _get_edge_info(self):
+    
+    def _get_edge_info_seastar(self):
         # getting the edge_list and edge_weights
         self._edge_list = []
         self._edge_weights = []
 
-        for time in range(self._dataset["time_periods"] - self.lags):
+        for time in range(self._dataset["time_periods"]):
             time_edge_list = []
             time_edge_weights = []
 
@@ -66,6 +70,18 @@ def _get_edge_info(self):
             self._edge_list.append(time_edge_list)
             self._edge_weights.append(time_edge_weights)
 
+    def _get_edge_info_pygt(self):
+        self._edge_list  = []
+        self._edge_weights = []
+        
+        for time in range(self._dataset["time_periods"]):
+            self._edge_list.append(
+                np.array(self._dataset["edge_mapping"]["edge_index"][str(time)]).T
+            )
+            self._edge_weights.append(
+                np.array(self._dataset["edge_mapping"]["edge_weight"][str(time)])
+            )
+
     def _get_targets_and_features(self):
         stacked_target = np.array(self._dataset["y"])
         standardized_target = (stacked_target - np.mean(stacked_target, axis=0)) / (
diff --git a/seastar/dataset/FoorahBase.py b/seastar/dataset/FoorahBase.py
new file mode 100644
index 0000000..a0e51b5
--- /dev/null
+++ b/seastar/dataset/FoorahBase.py
@@ -0,0 +1,160 @@
+import os
+import json
+import urllib.request
+import time
+
+import numpy as np
+
+from rich import inspect
+from rich.pretty import pprint
+from rich.progress import track
+from rich.console import Console
+
+console = Console()
+
+# from rich.traceback import install
+# install(show_locals=True)
+
+class FoorahBase:
+    def __init__(self, folder_name, dataset_name,verbose: bool = False, split=0.75, for_seastar= False) -> None:
+        self.name = dataset_name
+        self.split = split
+
+        self._graph_attr = {}
+        self._graph_updates = {}
+        self._max_num_nodes = 0
+
+        self._local_path = f'../../dataset/{folder_name}/{dataset_name}.json'
+        self._verbose = verbose
+
+        self._load_dataset()
+        self.total_timestamps = self._dataset["time_periods"]
+        self._get_targets_and_features()
+        
+        if for_seastar:
+            self._get_edge_info_seastar()
+            self._presort_edge_weights()
+        else:
+            self._get_edge_info_pygt()
+
+    def get_graph_data(self):
+        return self._graph_updates, self._max_num_nodes
+
+    def _load_dataset(self) -> None:
+        # loading the dataset by downloading them online
+        # if self._verbose:
+        #     console.log(f"Downloading [cyan]{self.name}[/cyan] dataset")
+        
+        # for online download
+        if os.path.exists(self._local_path):
+            dataset_file = open(self._local_path)
+            self._dataset = json.load(dataset_file)
+            if self._verbose:
+                console.log(f'Loading [cyan]{self.name}[/cyan] dataset from dataset/')
+        else:
+            console.log(f'Failed to find [cyan]{self.name}[/cyan] dataset from dataset')
+            quit()
+        
+        # for local
+        # dataset_file = open(self._local_path)
+        # self._dataset = json.load(dataset_file)
+
+    def _get_edge_info_seastar(self):
+        # getting the edge_list and edge_weights
+        self._edge_list = []
+        self._edge_weights = []
+
+        for time in range(self._dataset["time_periods"]):
+            time_edge_list = []
+            time_edge_weights = []
+
+            for edge in self._dataset["edge_mapping"]["edge_index"][str(time)]:
+                time_edge_list.append((edge[0], edge[1]))
+
+            for weight in self._dataset["edge_mapping"]["edge_weight"][str(time)]:
+                time_edge_weights.append(weight)
+
+            self._edge_list.append(time_edge_list)
+            self._edge_weights.append(time_edge_weights)
+    
+    def _get_edge_info_pygt(self):
+        self._edge_list  = []
+        self._edge_weights = []
+        
+        for time in range(self._dataset["time_periods"]):
+            self._edge_list.append(
+                np.array(self._dataset["edge_mapping"]["edge_index"][str(time)]).T
+            )
+            self._edge_weights.append(
+                np.array(self._dataset["edge_mapping"]["edge_weight"][str(time)])
+            )
+
+    def _get_targets_and_features(self):
+        stacked_target = np.array(self._dataset["y"])
+        standardized_target = (stacked_target - np.mean(stacked_target, axis=0)) / (
+            np.std(stacked_target, axis=0) + 10**-10
+        )
+
+        self._all_features = [
+            # standardized_target[i, :].T
+            standardized_target[i, :]
+            for i in range(self._dataset["time_periods"])
+        ]
+        self._all_targets = [
+            # standardized_target[i, :].T
+            standardized_target[i, :]
+            for i in range(self._dataset["time_periods"])
+        ]
+
+    def _presort_edge_weights(self):
+        """
+        Presorting edges according to (dest,src) since that is how eids are formed
+        allowing forward and backward kernel to access edge weights
+        """
+        final_edges_lst = []
+        final_edge_weights_lst = []
+
+        for i in range(len(self._edge_list)):
+            src_list = [edge[0] for edge in self._edge_list[i]]
+            dst_list = [edge[1] for edge in self._edge_list[i]]
+            weights = self._edge_weights[i]
+
+            edge_info_list = []
+            sorted_edges_lst = []
+            sorted_edge_weights_lst = []
+
+            for j in range(len(weights)):
+                edge_info = (src_list[j], dst_list[j], weights[j])
+                edge_info_list.append(edge_info)
+
+            # since it has to be sorted according to the reverse order
+            sorted_edge_info_list = sorted(
+                edge_info_list, key=lambda element: (element[1], element[0])
+            )
+
+            time_edge = []
+
+            for edge in sorted_edge_info_list:
+                time_edge.append((edge[0], edge[1]))
+                sorted_edge_weights_lst.append(edge[2])
+
+            final_edges_lst.append(time_edge)
+            final_edge_weights_lst.append(np.array(sorted_edge_weights_lst))
+
+        self._edge_list = final_edges_lst
+        self._edge_weights = final_edge_weights_lst
+        
+        # print("🍎🍎🍎 Edge Weights")
+        # print(self._edge_weights)
+
+    def get_edges(self):
+        return self._edge_list
+
+    def get_edge_weights(self):
+        return self._edge_weights
+
+    def get_all_features(self):
+        return self._all_features
+
+    def get_all_targets(self):
+        return self._all_targets
\ No newline at end of file
diff --git a/seastar/dataset/LinkPredDatasetLoader.py b/seastar/dataset/LinkPredDatasetLoader.py
new file mode 100644
index 0000000..bef6454
--- /dev/null
+++ b/seastar/dataset/LinkPredDatasetLoader.py
@@ -0,0 +1,90 @@
+import os
+import json
+import urllib.request
+import time
+
+import numpy as np
+import random
+
+from rich import inspect
+from rich.pretty import pprint
+from rich.progress import track
+from rich.console import Console
+import copy
+
+console = Console()
+
+# from rich.traceback import install
+# install(show_locals=True)
+
+class LinkPredDatasetLoader:
+    def __init__(self, folder_name, dataset_name, max_num_nodes, verbose: bool = False, for_seastar= False) -> None:
+        self.name = dataset_name
+
+        self._max_num_nodes = max_num_nodes
+
+        self._local_path = f'../../dataset/{folder_name}/{dataset_name}.json'
+        self._verbose = verbose
+
+        self._load_dataset()
+        self.total_timestamps = self._dataset["time_periods"]
+        self.for_seastar = for_seastar
+        
+        self._get_edge_info()
+        self._preprocess_pos_neg_edges()
+
+    def _load_dataset(self) -> None:
+        # loading the dataset by downloading them online
+        # if self._verbose:
+        #     console.log(f"Downloading [cyan]{self.name}[/cyan] dataset")
+        
+        # for online download
+        if os.path.exists(self._local_path):
+            dataset_file = open(self._local_path)
+            self._dataset = json.load(dataset_file)
+            if self._verbose:
+                console.log(f'Loading [cyan]{self.name}[/cyan] dataset from dataset/')
+        else:
+            console.log(f'Failed to find [cyan]{self.name}[/cyan] dataset from dataset')
+            quit()
+        
+        # for local
+        # dataset_file = open(self._local_path)
+        # self._dataset = json.load(dataset_file)
+
+    def _get_edge_info(self):
+        # getting the edge_list and edge_weights
+        edge_list = []
+        updates = self._dataset["edge_mapping"]["edge_index"]
+
+        working_set = set([(edge[0], edge[1]) for edge in updates["0"]["add"]])
+        edge_list.append(list(working_set))
+        for time in range(1, self.total_timestamps):
+            working_set = working_set.union(set([(edge[0], edge[1]) for edge in updates[str(time)]["add"]])).difference(set([(edge[0], edge[1]) for edge in updates[str(time)]["delete"]]))
+            edge_list.append(list(working_set))
+        
+        if self.for_seastar:
+            self._edge_list = edge_list
+        else:
+            self._edge_list = [np.array([[edge[0], edge[1]] for edge in edge_lst_t]).T for edge_lst_t in edge_list]
+    
+    def _preprocess_pos_neg_edges(self):
+        updates = self._dataset["edge_mapping"]["edge_index"]
+
+        pos_neg_edge_list = []
+        pos_neg_edge_label_list = []
+
+        for i in range(1, self.total_timestamps):
+            pos_edges_tup = list(updates[str(i)]["add"])
+            neg_edges_tup = list(updates[str(i)]["neg"])
+            pos_neg_edge_list.append(pos_edges_tup + neg_edges_tup)
+            pos_neg_edge_label_list.append([(edge[0], edge[1], 1) for edge in pos_edges_tup] + [(edge[0], edge[1], 0) for edge in neg_edges_tup])
+
+        self._pos_neg_edge_list = [np.array(edge_list).T for edge_list in pos_neg_edge_list]
+        self._pos_neg_edge_label_list = [np.array([edge[2] for edge in edge_list]) for edge_list in pos_neg_edge_label_list]
+
+    def get_edges(self):
+        return self._edge_list
+    
+    def get_pos_neg_edges(self):
+        return self._pos_neg_edge_list, self._pos_neg_edge_label_list
\ No newline at end of file
diff --git a/seastar/dataset/SoorahBase.py b/seastar/dataset/SoorahBase.py
index 73dd06d..a86ca38 100644
--- a/seastar/dataset/SoorahBase.py
+++ b/seastar/dataset/SoorahBase.py
@@ -145,7 +145,7 @@ def _presort_edge_weights(self):
         self._edge_list = final_edges_lst
         self._edge_weights = final_edge_weights_lst
         
-        print("🍎🍎🍎 Edge Weights")
+        # print("🍎🍎🍎 Edge Weights")
         # print(self._edge_weights)
 
     def get_edges(self):
diff --git a/seastar/graph/SeastarGraph.py b/seastar/graph/SeastarGraph.py
index 5edd24e..0833192 100644
--- a/seastar/graph/SeastarGraph.py
+++ b/seastar/graph/SeastarGraph.py
@@ -2,7 +2,7 @@
 
 class SeastarGraph(ABC):
     def __init__(self):
-        self.ndata = {}
+        self._ndata = {}
         
         self._forward_graph = None
         self._backward_graph = None
@@ -26,6 +26,14 @@ def get_num_nodes(self):
     @abstractmethod
     def get_num_edges(self):
         pass
+
+    @abstractmethod
+    def get_ndata(self, field):
+        pass
+
+    @abstractmethod
+    def set_ndata(self, field, val):
+        pass
     
     @property
     @abstractmethod
diff --git a/seastar/graph/dynamic/DynamicGraph.py b/seastar/graph/dynamic/DynamicGraph.py
index 3f82ed8..34238a7 100644
--- a/seastar/graph/dynamic/DynamicGraph.py
+++ b/seastar/graph/dynamic/DynamicGraph.py
@@ -1,62 +1,20 @@
-import copy
-import time
-
-
 from seastar.graph.SeastarGraph import SeastarGraph
-from rich import inspect
-
-
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 
 class DynamicGraph(SeastarGraph):
-    def __init__(self, edge_list):
+    def __init__(self, edge_list, max_num_nodes):
         super().__init__()
         self.graph_updates = {}
-        self.max_num_nodes = 0
-        self.graph_cache = {}
-        
+        self.max_num_nodes = max_num_nodes
+        self.graph_attr = {str(t): (self.max_num_nodes, len(set(edge_list[t]))) for t in range(len(edge_list))}
+
         # Indicates whether the graph is currently undergoing backprop
         self._is_backprop_state = False
         self.current_timestamp = 0
-        
+
         self._preprocess_graph_structure(edge_list)
-        
-    def _get_graph_attr(self, edge_list):
-        graph_attr = {}
-        for time in range(len(edge_list)):
-            node_set = set()
-            edge_count = 0
-            for edge in edge_list[time]:
-                src = edge[0]
-                dst = edge[1]
-                edge_count += 1
-                node_set.add(src)
-                node_set.add(dst)
-            node_count = len(node_set)
-            graph_attr[str(time)] = (node_count, edge_count)
-        
-        return graph_attr
-        
-    def _preprocess_graph_structure(self, edge_list):
-        
-        graph_attr = self._get_graph_attr(edge_list)
-        
-        max_num_nodes = 0
-        for i in range(len(edge_list)):
-            for j in range(len(edge_list[i])):
-                max_num_nodes = max(max_num_nodes,edge_list[i][j][0],edge_list[i][j][1])
-        self.max_num_nodes = max_num_nodes + 1
-        
-        # SINCE THE CONCEPT OF NODE IDS HASNT BEEN IMPLEMENTED
-        # WE CANT USE THIS (AS in consider the case where nodes are labelled 1,2,3)
-        # What hapens to 0, Our system right now cant map each node to an ID
-        # tmp_set = set()
-        # for i in range(len(edge_list)):
-        #     for j in range(len(edge_list[i])):
-        #         tmp_set.add(edge_list[i][j][0])
-        #         tmp_set.add(edge_list[i][j][1])
-        # self.max_num_nodes = len(tmp_set)
 
+    def _preprocess_graph_structure(self, edge_list):
         edge_dict = {}
         for i in range(len(edge_list)):
             edge_set = set()
@@ -65,71 +23,73 @@ def _preprocess_graph_structure(self, edge_list):
             edge_dict[str(i)] = edge_set
 
         self.graph_updates = {}
-        self.graph_updates["0"] = {
-            "add": list(edge_dict["0"]),
-            "delete": [],
-            "num_nodes": graph_attr["0"][0],
-            "num_edges": graph_attr["0"][1],
-        }
+        self.graph_updates["0"] = {"add": list(edge_dict["0"]), "delete": []}
         for i in range(1, len(edge_list)):
             self.graph_updates[str(i)] = {
                 "add": list(edge_dict[str(i)].difference(edge_dict[str(i - 1)])),
                 "delete": list(edge_dict[str(i - 1)].difference(edge_dict[str(i)])),
-                "num_nodes": graph_attr[str(i)][0],
-                "num_edges": graph_attr[str(i)][1],
             }
-        
+
     def get_graph(self, timestamp: int):
-        # print("💄💄💄 Get_graph (forward) called",flush=True)
         self._is_backprop_state = False
-        
+
         if timestamp < self.current_timestamp:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_forward()")
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_forward()"
+            )
 
         while self.current_timestamp < timestamp:
             self._update_graph_forward()
             self.current_timestamp += 1
 
     def get_backward_graph(self, timestamp: int):
-        # print("📝📝📝 Get_backward_graph (backward) called",flush=True)
         if not self._is_backprop_state:
-            # print("💄💄💄 Calling backward init")
             self._is_backprop_state = True
             self._init_reverse_graph()
-        
+
         if timestamp > self.current_timestamp:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_backward()")
-        
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_backward()"
+            )
+
         while self.current_timestamp > timestamp:
-            # print("🎒🎒🎒 Calling update backward")
             self._update_graph_backward()
             self.current_timestamp -= 1
-    
+
     def get_num_nodes(self):
-        return self.graph_updates[str(self.current_timestamp)]["num_nodes"]
-    
+        return self.graph_attr[str(self.current_timestamp)][0]
+
     def get_num_edges(self):
-        return self.graph_updates[str(self.current_timestamp)]["num_edges"]
+        return self.graph_attr[str(self.current_timestamp)][1]
     
+    def get_ndata(self, field):
+        if str(self.current_timestamp) in self._ndata and field in self._ndata[str(self.current_timestamp)]:
+            return self._ndata[str(self.current_timestamp)][field]
+        else:
+            return None
+
+    def set_ndata(self, field, val):
+        if str(self.current_timestamp) in self._ndata:
+            self._ndata[str(self.current_timestamp)][field] = val
+        else:
+            self._ndata[str(self.current_timestamp)] = {field: val}
+
     @abstractmethod
     def in_degrees(self):
         pass
-    
+
     @abstractmethod
     def out_degrees(self):
         pass
-    
-    
+
     @abstractmethod
     def _update_graph_forward(self):
         pass
-    
+
     @abstractmethod
     def _init_reverse_graph(self):
         pass
-    
+
     @abstractmethod
     def _update_graph_backward(self):
         pass
-    
-    
\ No newline at end of file
diff --git a/seastar/graph/dynamic/gpma/GPMAGraph.py b/seastar/graph/dynamic/gpma/GPMAGraph.py
index 4133437..e2d9b35 100644
--- a/seastar/graph/dynamic/gpma/GPMAGraph.py
+++ b/seastar/graph/dynamic/gpma/GPMAGraph.py
@@ -1,121 +1,130 @@
 import copy
+import time
+
 import numpy as np
 from rich import inspect
 
 from seastar.graph.dynamic.DynamicGraph import DynamicGraph
-from seastar.graph.dynamic.gpma.gpma import GPMA, init_gpma, print_gpma_info, edge_update_list, label_edges, copy_label_edges, build_reverse_gpma, get_csr_ptrs, get_graph_attr
+from seastar.graph.dynamic.gpma.gpma import (
+    GPMA,
+    init_gpma,
+    init_graph_updates,
+    edge_update_t,
+    label_edges,
+    build_backward_csr,
+    free_backward_csr,
+    get_csr_ptrs,
+    get_out_degrees,
+    get_in_degrees
+)
 
 
 class GPMAGraph(DynamicGraph):
-    def __init__(self, edge_list):
-        super().__init__(edge_list)
-        
-        # forward and backward graphs for GPMA
+    def __init__(self, edge_list, max_num_nodes):
+        super().__init__(edge_list, max_num_nodes)
+
+        # forward graph for GPMA
         self._forward_graph = GPMA()
-        self._backward_graph = GPMA()
-        
         init_gpma(self._forward_graph, self.max_num_nodes)
-        init_gpma(self._backward_graph, self.max_num_nodes)
-        
-        # initialise the graph for the first time stamp
-        initial_graph_additions = self.graph_updates["0"]["add"]
+        init_graph_updates(self._forward_graph, self.graph_updates, reverse_edges=True)
 
-        edge_update_list(self._forward_graph, initial_graph_additions, is_reverse_edge=True)
+        # base forward graph at t=0
+        edge_update_t(self._forward_graph, 0)
         label_edges(self._forward_graph)
 
+        # Cacheing node degrees
+        self._in_degrees_cache = {}
+        self._out_degrees_cache = {}
+
+        # for benchmarking purposes
+        self._update_count = 0
+        self._total_update_time = 0
+        self._gpu_move_time = 0
+
         self._get_graph_csr_ptrs()
-        
+
     def graph_type(self):
         return "gpma"
-        
+    
     def in_degrees(self):
-        return np.array(self._forward_graph.out_degree, dtype='int32')
+        return np.array(get_out_degrees(self._forward_graph), dtype="int32")
     
     def out_degrees(self):
-        return np.array(self._forward_graph.in_degree, dtype='int32')
-    
-    def _update_reverse_graph_cache(self):
-        # saving reverse base graph in cache
-        self.graph_cache['reverse'] = copy.deepcopy(self._backward_graph)
-            
-    def _get_reverse_cached_graph(self):
-        return copy.deepcopy(self.graph_cache['reverse'])
-    
+        return np.array(get_in_degrees(self._forward_graph), dtype="int32")
+
     def _get_graph_csr_ptrs(self):
+
         forward_csr_ptrs = get_csr_ptrs(self._forward_graph)
-        backward_csr_ptrs = get_csr_ptrs(self._backward_graph)
-        
         self.fwd_row_offset_ptr = forward_csr_ptrs[0]
         self.fwd_column_indices_ptr = forward_csr_ptrs[1]
         self.fwd_eids_ptr = forward_csr_ptrs[2]
-        
-        self.bwd_row_offset_ptr = backward_csr_ptrs[0]
-        self.bwd_column_indices_ptr = backward_csr_ptrs[1]
-        self.bwd_eids_ptr = backward_csr_ptrs[2]
-    
-    #TODO: Right now this returns (max_num_nodes,num_edges) see if this is what is required
-    # def _get_graph_attributes(self):
-
-    #     if not self._is_backprop_state:
-    #         graph_attr = get_graph_attr(self._forward_graph)
-    #     else:
-    #         graph_attr = get_graph_attr(self._backward_graph)
-        
-    #     self.num_nodes = graph_attr[0]
-    #     self.num_edges = graph_attr[1]
-    
+
+        if self._is_backprop_state:
+          backward_csr_ptrs = get_csr_ptrs(self._forward_graph, is_backward=True)
+          self.bwd_row_offset_ptr = backward_csr_ptrs[0]
+          self.bwd_column_indices_ptr = backward_csr_ptrs[1]
+          self.bwd_eids_ptr = backward_csr_ptrs[2]
+
     def _update_graph_forward(self):
-        
         # if we went through the entire time-stamps
         if str(self.current_timestamp + 1) not in self.graph_updates:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_forward()")
-        
-        # getting the graph edge modifications in 
-        # the following form list[tuple(int, int)]
-        graph_additions = self.graph_updates[str(self.current_timestamp + 1)]["add"]
-        graph_deletions = self.graph_updates[str(self.current_timestamp + 1)]["delete"]
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_forward()"
+            )
 
-        edge_update_list(self._forward_graph, graph_additions, is_reverse_edge=True)
-        edge_update_list(self._forward_graph, graph_deletions, is_delete=True, is_reverse_edge=True)
+        self._update_count += len(
+            self.graph_updates[str(self.current_timestamp + 1)]["add"]
+        )
+        self._update_count += len(
+            self.graph_updates[str(self.current_timestamp + 1)]["delete"]
+        )
 
-        # TODO: UNCOMMENT LATER
+        update_time_0 = time.time()
+
+        edge_update_t(self._forward_graph, self.current_timestamp + 1)
         label_edges(self._forward_graph)
+
+        update_time_1 = time.time()
+        self._total_update_time += update_time_1 - update_time_0
+
         self._get_graph_csr_ptrs()
-        # self._get_graph_attributes()  # NOTE:
-        
+
     def _init_reverse_graph(self):
-        ''' Generates the reverse of the base graph'''
+        """Generates the reverse of the base graph"""
+        update_time_0 = time.time()
 
-        # checking if the reverse base graph exists in the cache
-        # we can load it from there instead of building it each time
-        if 'reverse' in self.graph_cache:
-            self._backward_graph = self._get_reverse_cached_graph()
-        else:
-            build_reverse_gpma(self._backward_graph, self._forward_graph)
+        free_backward_csr(self._forward_graph)
+        build_backward_csr(self._forward_graph)
+        self._get_graph_csr_ptrs()
 
-            # storing the reverse base graph in cache after building
-            # it for the first time
-            self._update_reverse_graph_cache()
+        update_time_1 = time.time()
+        self._total_update_time += update_time_1 - update_time_0
 
-        self._get_graph_csr_ptrs()
-        # self._get_graph_attributes() # NOTE:
-        
     def _update_graph_backward(self):
         if self.current_timestamp < 0:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_backward()")
-        
-        graph_additions = self.graph_updates[str(self.current_timestamp)]["delete"]
-        graph_deletions = self.graph_updates[str(self.current_timestamp)]["add"]
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_backward()"
+            )
 
-        edge_update_list(self._backward_graph, graph_additions)
-        edge_update_list(self._backward_graph, graph_deletions, is_delete=True)
+        self._update_count += len(
+            self.graph_updates[str(self.current_timestamp)]["delete"]
+        )
+        self._update_count += len(
+            self.graph_updates[str(self.current_timestamp)]["add"]
+        )
 
-        edge_update_list(self._forward_graph, graph_additions, is_reverse_edge=True)
-        edge_update_list(self._forward_graph, graph_deletions, is_delete=True, is_reverse_edge=True)
+        update_time_0 = time.time()
 
-        # TODO: UNCOMMENT LATER
+        # Freeing resources from previous CSR
+        free_backward_csr(self._forward_graph)
+
+        edge_update_t(
+            self._forward_graph, self.current_timestamp, revert_update=True
+        )
         label_edges(self._forward_graph)
-        copy_label_edges(self._backward_graph, self._forward_graph)
+        build_backward_csr(self._forward_graph)
+
+        update_time_1 = time.time()
+        self._total_update_time += update_time_1 - update_time_0
 
         self._get_graph_csr_ptrs()
-        # self._get_graph_attributes()  # NOTE:
\ No newline at end of file
diff --git a/seastar/graph/dynamic/gpma/gpma.cu b/seastar/graph/dynamic/gpma/gpma.cu
index 77ff211..8cfef07 100644
--- a/seastar/graph/dynamic/gpma/gpma.cu
+++ b/seastar/graph/dynamic/gpma/gpma.cu
@@ -12,7 +12,7 @@
 #include <string>
 #include <vector>
 #include <tuple>
-// #include <memory>
+#include <chrono>
 #include "stdio.h"
 
 namespace py = pybind11;
@@ -55,30 +55,46 @@ const SIZE_TYPE MAX_BLOCKS_NUM = 96 * 8;
 class GPMA
 {
 public:
-    // NOTE: Assuming that keys holds the column_indices
-    // and values is the column indices (subject to change)
-    DEV_VEC_KEY keys;
-    DEV_VEC_VALUE values;
+    DEV_VEC_SIZE row_offset; // row offset vector
+    DEV_VEC_KEY keys;   // column indices vector
+    DEV_VEC_VALUE values;   // edge IDs vector
 
+    // node and edge metadata
+    SIZE_TYPE row_num; // number of nodes
+    SIZE_TYPE edge_count = 0;
+
+    // node degree vectors
+    DEV_VEC_SIZE in_degree;
+    DEV_VEC_SIZE out_degree;
+    DEV_VEC_SIZE cum_out_degree; // cummulative out degree
+
+    // pointers to pinned node degree vectors
+    SIZE_TYPE* in_degree_host;
+    SIZE_TYPE* out_degree_host;
+
+    // pointers to reverse CSR
+    SIZE_TYPE* bwd_row_offset = NULL;
+    KEY_TYPE* bwd_keys = NULL;
+    VALUE_TYPE* bwd_values = NULL;
+
+    // (add|delete)_updates stores the add/delete updates of every timestamp
+    // (add|delete)_updates_count maintains the number of additions/deletions
+    // in every timestamp
+    std::vector<KEY_TYPE *> add_updates;
+    std::vector<int> add_updates_count;
+    std::vector<KEY_TYPE *> delete_updates;
+    std::vector<int> delete_updates_count;
+
+    // GPMA helper datastructures
     SIZE_TYPE segment_length;
     SIZE_TYPE tree_height;
-
     double density_lower_thres_leaf = 0.08;
     double density_lower_thres_root = 0.42;
     double density_upper_thres_root = 0.84;
     double density_upper_thres_leaf = 0.92;
-
     thrust::host_vector<SIZE_TYPE> lower_element;
     thrust::host_vector<SIZE_TYPE> upper_element;
 
-    // addition for csr
-    SIZE_TYPE row_num; // number of nodes
-    DEV_VEC_SIZE row_offset;
-
-    int edge_count = 0;
-    std::vector<unsigned int> in_degree;
-    std::vector<unsigned int> out_degree;
-
     GPMA();
 
     inline int get_size()
@@ -161,23 +177,25 @@ __host__ void recalculate_density(GPMA &gpma)
 __device__ void cub_sort_key_value(KEY_TYPE *keys, VALUE_TYPE *values, SIZE_TYPE size, KEY_TYPE *tmp_keys,
                                    VALUE_TYPE *tmp_values)
 {
+    // Create a set of DoubleBuffers to wrap pairs of device pointers
+    cub::DoubleBuffer<KEY_TYPE> d_keys(keys, tmp_keys);
+    cub::DoubleBuffer<VALUE_TYPE> d_values(values, tmp_values);
 
     void *d_temp_storage = NULL;
     size_t temp_storage_bytes = 0;
 
-    cErr(cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, keys, tmp_keys, values, tmp_values, size));
+    cErr(cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, size));
     cErr(cudaDeviceSynchronize());
     cErr(cudaMalloc(&d_temp_storage, temp_storage_bytes));
     cErr(cudaDeviceSynchronize());
-    cErr(cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, keys, tmp_keys, values, tmp_values, size));
+    cErr(cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, size));
     cErr(cudaDeviceSynchronize());
 
     SIZE_TYPE THREADS_NUM = 128;
     SIZE_TYPE BLOCKS_NUM = CALC_BLOCKS_NUM(THREADS_NUM, size);
-    memcpy_kernel<KEY_TYPE><<<BLOCKS_NUM, THREADS_NUM>>>(keys, tmp_keys, size);
-    memcpy_kernel<VALUE_TYPE><<<BLOCKS_NUM, THREADS_NUM>>>(values, tmp_values, size);
+    memcpy_kernel<KEY_TYPE><<<BLOCKS_NUM, THREADS_NUM>>>(d_keys.Alternate(), d_keys.Current(), size);
+    memcpy_kernel<VALUE_TYPE><<<BLOCKS_NUM, THREADS_NUM>>>(d_values.Alternate(), d_values.Current(), size);
     cErr(cudaDeviceSynchronize());
-
     cErr(cudaFree(d_temp_storage));
 }
 
@@ -923,10 +941,15 @@ __host__ void init_gpma(GPMA &gpma, SIZE_TYPE row_num)
     gpma.row_num = row_num;
     gpma.row_offset.resize(row_num + 1, 0);
 
-    // initialising in_degree and out_degree arrays
+    // initialising in_degree, out_degree and cum_out_degree arrays
     // with all zero values
     gpma.in_degree.resize(row_num, 0);
     gpma.out_degree.resize(row_num, 0);
+    gpma.cum_out_degree.resize(row_num, 0);
+
+    // allocating pinned memory for node degree vectors
+    cErr(cudaMallocHost(&gpma.in_degree_host, sizeof(SIZE_TYPE) * row_num));
+    cErr(cudaMallocHost(&gpma.out_degree_host, sizeof(SIZE_TYPE) * row_num));
 
     // creates a device vector of size row_num
     DEV_VEC_KEY row_wall(row_num);
@@ -943,395 +966,466 @@ __host__ void init_gpma(GPMA &gpma, SIZE_TYPE row_num)
     cErr(cudaDeviceSynchronize());
 }
 
-///////////////////////////////////////////////////////////////////////////////////////////////////
+//////////////////////////////////////// GPMA APIs ///////////////////////////////////////////////////////////
 
-void print_gpma_info(GPMA &gpma, int node)
+void init_graph_updates(GPMA &gpma, std::map<std::string, std::map<std::string, std::vector<std::tuple<int, int>>>> updates, bool reverse_edges = false)
 {
-    // Pretty prints the row_offset range and column_indices values
-    // for the given node
-    //
-    // row_offset:  Row offset vector
-    // col_indices: Column indices PMA
-    // node:        Node whose row_offset and column_indices is to be printed
 
-    // getting the (beg,end) range of the node
+    gpma.add_updates.resize(updates.size());
+    gpma.delete_updates.resize(updates.size());
+    gpma.add_updates_count.resize(updates.size());
+    gpma.delete_updates_count.resize(updates.size());
 
-    thrust::host_vector<SIZE_TYPE> row_offset = gpma.row_offset;
-    thrust::host_vector<KEY_TYPE> col_indices = gpma.keys;
-    thrust::host_vector<VALUE_TYPE> edge_values = gpma.values;
+    std::vector<std::tuple<int, int>> update_tup;
 
-    unsigned int beg = row_offset[node];
-    unsigned int end = row_offset[node + 1];
-    int print_count = 0;
-
-    std::cout << "\n🌟 Column Indices for node: " << node << "(" << beg << ", " << end << ")"
-              << "\n\n";
-    std::cout << std::setw(6);
-
-    for (int i = beg; i < end; ++i)
+    for (int t = 0; t < updates.size(); ++t)
     {
-        KEY_TYPE mask = (KEY_TYPE)node << 32;
-        unsigned int dst = (col_indices[i] - mask);
-        VALUE_TYPE val = edge_values[i];
-
-        if (dst != COL_IDX_NONE && val != VALUE_NONE)
+        // preparing keys of edges to be added
+        update_tup = updates[std::to_string(t)]["add"];
+        std::vector<KEY_TYPE> add_key(update_tup.size());
+        for (int i = 0; i < update_tup.size(); ++i)
         {
-            std::cout << dst << "(" << val << ")" << std::setw(6);
+            unsigned int src = reverse_edges ? std::get<1>(update_tup[i]) : std::get<0>(update_tup[i]);
+            unsigned int dst = reverse_edges ? std::get<0>(update_tup[i]) : std::get<1>(update_tup[i]);
+            add_key[i] = ((KEY_TYPE)src << 32) + dst;
         }
-        else
+
+        // preparing keys of edges to be deleted
+        update_tup = updates[std::to_string(t)]["delete"];
+        std::vector<KEY_TYPE> delete_key(update_tup.size());
+        for (int i = 0; i < update_tup.size(); ++i)
         {
-            std::cout << "-" << std::setw(6);
+            unsigned int src = reverse_edges ? std::get<1>(update_tup[i]) : std::get<0>(update_tup[i]);
+            unsigned int dst = reverse_edges ? std::get<0>(update_tup[i]) : std::get<1>(update_tup[i]);
+            delete_key[i] = ((KEY_TYPE)src << 32) + dst;
         }
-        print_count += 1;
 
-        if (print_count > 10)
-        {
-            printf("\n");
-            print_count = 0;
+        // allocating and copying to pinned memory
+        KEY_TYPE *pinned_add_key;
+        KEY_TYPE *pinned_delete_key;
+        cErr(cudaMallocHost(&pinned_add_key, sizeof(KEY_TYPE) * add_key.size()));
+        cErr(cudaMallocHost(&pinned_delete_key, sizeof(KEY_TYPE) * delete_key.size()));
+        cErr(cudaMemcpy(pinned_add_key, add_key.data(), sizeof(KEY_TYPE) * add_key.size(), cudaMemcpyHostToHost));
+        cErr(cudaMemcpy(pinned_delete_key, delete_key.data(), sizeof(KEY_TYPE) * delete_key.size(), cudaMemcpyHostToHost));
+
+        // storing the pointers and counts in GPMA object
+        gpma.add_updates[t] = pinned_add_key;
+        gpma.delete_updates[t] = pinned_delete_key;
+        gpma.add_updates_count[t] = updates[std::to_string(t)]["add"].size();
+        gpma.delete_updates_count[t] = updates[std::to_string(t)]["delete"].size();
+    }
+}
+
+__global__ void update_node_degrees_kernel(SIZE_TYPE* in_degree, SIZE_TYPE* out_degree, KEY_TYPE* updates, int updates_size, bool is_delete){
+    int index = blockDim.x * blockIdx.x + threadIdx.x;
+    int block_offset = gridDim.x * blockDim.x;
+    for(; index < updates_size; index += block_offset){
+        KEY_TYPE key = updates[index];
+        SIZE_TYPE src = (SIZE_TYPE)(key >> 32);
+        SIZE_TYPE dst = (SIZE_TYPE)(key);
+
+        if(is_delete){
+            atomicSub(&in_degree[dst],1);
+            atomicSub(&out_degree[src],1);
+        }else{
+            atomicAdd(&in_degree[dst],1);
+            atomicAdd(&out_degree[src],1);
         }
     }
-    printf("\n\n");
 }
 
-void load_data(const char *file_path, thrust::host_vector<int> &host_x, thrust::host_vector<int> &host_y,
-               int &node_size, int &edge_size)
-{
+void update_node_degrees(SIZE_TYPE* in_degree, SIZE_TYPE* out_degree, KEY_TYPE* add_keys, int add_key_size, KEY_TYPE* del_keys, int del_key_size){
+    
+    // Updating node degrees associated with added edges
+    SIZE_TYPE THREADS_NUM = 128;
+    SIZE_TYPE BLOCKS_NUM = CALC_BLOCKS_NUM(THREADS_NUM, (SIZE_TYPE)add_key_size);
+    update_node_degrees_kernel<<<BLOCKS_NUM,THREADS_NUM>>>(in_degree, out_degree, add_keys, add_key_size,false);
 
-    FILE *fp;
-    fp = fopen(file_path, "r");
-    if (not fp)
-    {
-        printf("Open graph file failed.\n");
-        exit(0);
-    }
+    // Updating node degrees associated with deleted edges
+    BLOCKS_NUM = CALC_BLOCKS_NUM(THREADS_NUM, (SIZE_TYPE)del_key_size);
+    update_node_degrees_kernel<<<BLOCKS_NUM,THREADS_NUM>>>(in_degree, out_degree, del_keys, del_key_size, true);
+    cErr(cudaDeviceSynchronize());
+}
 
-    // gets the number of nodes and edges from the first line
-    // [num_nodes] [num_edges]
-    fscanf(fp, "%d %d", &node_size, &edge_size);
-    printf("node_num: %d, edge_num: %d\n", node_size, edge_size);
+std::vector<float> edge_update_t(GPMA &gpma, int timestamp, bool revert_update = false)
+{   
+    // This function performs/reverts the updates of a given timestamp
+    // based on if revert_update is false or true respectively.
 
-    // dynamic resizing to fit all edges
-    host_x.resize(edge_size);
-    host_y.resize(edge_size);
+    int add_edge_count, delete_edge_count;
+    KEY_TYPE *add_updates_ptr, *delete_updates_ptr;
+    KEY_TYPE *add_key_device, *delete_key_device;
 
-    // loads the src and dst nodes into x_host
-    // and y_host respectively
-    for (int i = 0; i < edge_size; i++)
+    if (revert_update)
     {
-        int x, y;
-
-        // [src_node] [dst_node]
-        fscanf(fp, "%d %d", &x, &y);
-        host_x[i] = x;
-        host_y[i] = y;
+        add_edge_count = gpma.delete_updates_count[timestamp];
+        delete_edge_count = gpma.add_updates_count[timestamp];
+        add_updates_ptr = gpma.delete_updates[timestamp];
+        delete_updates_ptr = gpma.add_updates[timestamp];
     }
-
-    printf("Graph file is loaded.\n");
-    fclose(fp);
-}
-
-void load_graph(GPMA &gpma, const char *file_path)
-{
-    // sets the GPU Malloc Heap Size to 1GB
-    // we could change the limit accordingly
-    // cudaDeviceSetLimit(cudaLimitMallocHeapSize, 1024ll * 1024);
-
-    // controls the maximum nesting depth of a grid at which
-    // a thread can safely call cudaDeviceSynchronize()
-    // cudaDeviceSetLimit(cudaLimitDevRuntimeSyncDepth, 5);
-
-    // we are defining two host_vectors on the CPU
-    // host_x will contain the src nodes and
-    // host_y will contain the destination nodes of the edge
-    thrust::host_vector<int> host_x;
-    thrust::host_vector<int> host_y;
-
-    int node_size; // num of nodes
-    int edge_size; // num of edges
-
-    // loads the graph edges
-    load_data(file_path, host_x, host_y, node_size, edge_size);
-
-    // we are starting out with half the number of edges
-    // as the initial graph. After this is done, the other
-    // half of edges are added dynamically to the graph in
-    // batches of 100 edges
-    int half = edge_size / 2;
-
-    // h_base_key is a host_vector to maintain a unique ID
-    // for each edge. KEY_TYPE = unsigned long long.
-    // The unique ID is a 64-bit number, where the first 32 bits
-    // represent the source node and the last 32-bits the
-    // destination node.
-    thrust::host_vector<KEY_TYPE> h_base_keys(half);
-
-    // looping over half the edges to create the unique edge ID.
-    // It shifts the source node ID (64-bit after casting) by 32 bits
-    // to the left and the 32-bit destination node ID is added
-    // to the 32-bits on the left of source node ID
-    for (int i = 0; i < half; i++)
-        h_base_keys[i] = ((KEY_TYPE)host_x[i] << 32) + host_y[i];
-
-    // device_vector<unsigned long long>. The unique edge IDs are
-    // copied into the GPU device through a device_vector
-    DEV_VEC_KEY base_keys = h_base_keys;
-
-    // device_vector<double>. device_vector of size
-    // half = num_edges/2. Contains the value of the edges.
-    // Initialised to 1 and stored in GPU device
-    DEV_VEC_VALUE base_values(half, 1);
-
-    // the base keys and values are being transferred to
-    // GPU device. Host is waiting for that task to complete
-    // before continuing with the rest.
+    else
+    {
+        add_edge_count = gpma.add_updates_count[timestamp];
+        delete_edge_count = gpma.delete_updates_count[timestamp];
+        add_updates_ptr = gpma.add_updates[timestamp];
+        delete_updates_ptr = gpma.delete_updates[timestamp];
+    }
+    gpma.edge_count = gpma.edge_count + add_edge_count - delete_edge_count;
+
+    // Preparing data for addition updates
+    cErr(cudaMalloc(&add_key_device, sizeof(KEY_TYPE) * add_edge_count));
+    cErr(cudaMemcpy(add_key_device, add_updates_ptr, sizeof(KEY_TYPE) * add_edge_count, cudaMemcpyHostToDevice));
+    thrust::device_ptr<KEY_TYPE> add_key_thrust_dev_ptr = thrust::device_pointer_cast(add_key_device);
+    DEV_VEC_KEY add_key_thrust_dev(add_key_thrust_dev_ptr, add_key_thrust_dev_ptr + add_edge_count);
+    thrust::device_vector<VALUE_TYPE> add_value_thrust_dev(add_edge_count, 1);
     cudaDeviceSynchronize();
 
-    init_gpma(gpma, node_size); //
+    // Preparing data for deletion updates
+    cErr(cudaMalloc(&delete_key_device, sizeof(KEY_TYPE) * delete_edge_count));
+    cErr(cudaMemcpy(delete_key_device, delete_updates_ptr, sizeof(KEY_TYPE) * delete_edge_count, cudaMemcpyHostToDevice));
+    thrust::device_ptr<KEY_TYPE> delete_key_thrust_dev_ptr = thrust::device_pointer_cast(delete_key_device);
+    DEV_VEC_KEY delete_key_thrust_dev(delete_key_thrust_dev_ptr, delete_key_thrust_dev_ptr + delete_edge_count);
+    thrust::device_vector<VALUE_TYPE> delete_value_thrust_dev(delete_edge_count, VALUE_NONE);
     cudaDeviceSynchronize();
 
-    update_gpma(gpma, base_keys, base_values);
+    // Performing GPMA update
+    auto start_time_update = std::chrono::high_resolution_clock::now();
+    update_gpma(gpma, add_key_thrust_dev, add_value_thrust_dev);
     cudaDeviceSynchronize();
-
-    printf("Graph is updated.\n");
+    update_gpma(gpma, delete_key_thrust_dev, delete_value_thrust_dev);
+    cudaDeviceSynchronize();
+    auto end_time_update = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<float> time_update = (end_time_update - start_time_update);
+
+    // Performing node degree update
+    auto start_time_node_degrees = std::chrono::high_resolution_clock::now();
+    update_node_degrees(RAW_PTR(gpma.in_degree), RAW_PTR(gpma.out_degree), add_key_device, add_edge_count, delete_key_device, delete_edge_count);
+    auto end_time_node_degrees = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<float> time_node_degree = (end_time_node_degrees - start_time_node_degrees);
+
+    // freeing resources
+    cErr(cudaFree(add_key_device));
+    cErr(cudaFree(delete_key_device));
+
+    std::vector<float> vec;
+    vec.push_back(time_update.count());
+    vec.push_back(time_node_degree.count());
+    return vec;
 }
 
-void edge_update_list(GPMA &gpma, std::vector<std::tuple<int, int>> edge_list, bool is_delete = false, bool is_reverse_edge = false)
-{
-    // NOTE:: Should we set these limits every single time?
-    // cudaDeviceSetLimit(cudaLimitMallocHeapSize, 1024ll * 1024);
-    // cudaDeviceSetLimit(cudaLimitDevRuntimeSyncDepth, 5);
-
-    thrust::host_vector<int> host_src;
-    thrust::host_vector<int> host_dst;
+__global__ void label_edges_kernel(SIZE_TYPE* row_offset, KEY_TYPE* column_indices, VALUE_TYPE* eids, SIZE_TYPE* cum_out_degree, int cum_out_degree_size){
+    
+    int index = blockDim.x * blockIdx.x + threadIdx.x;
+    int block_offset = gridDim.x * blockDim.x;
+    int edge_count = 1;
+    int beg, end;
 
-    int edge_count = 0;
+    for(; index < cum_out_degree_size; index += block_offset){
+        beg = row_offset[index];
+        end = row_offset[index+1];
 
-    // iterating through a vector of tuples
-    // each tuple is of the form (src_node, dst_node)
+        if(index > 0){
+            edge_count = cum_out_degree[index-1]+1;
+        }
+        
+        for(int i=beg; i<end; ++i){
+            KEY_TYPE key = column_indices[i];
+            VALUE_TYPE value = eids[i];
+            if (key != KEY_MAX && (key & 0xffffffff) != COL_IDX_NONE && value != VALUE_NONE)
+            {
+                eids[i] = edge_count;
+                ++edge_count;
+            }
+        }
+    }
+}
 
-    for (auto &edge : edge_list)
-    {
+void label_edges(GPMA &gpma){
+    // Calculating cummulative out degrees
+    void *d_temp_storage = NULL;
+    size_t temp_storage_bytes = 0;
+    cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, RAW_PTR(gpma.out_degree), RAW_PTR(gpma.cum_out_degree), gpma.row_num);
+    cudaMalloc(&d_temp_storage, temp_storage_bytes);
+    cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, RAW_PTR(gpma.out_degree), RAW_PTR(gpma.cum_out_degree), gpma.row_num);
+    cErr(cudaDeviceSynchronize());
+    cErr(cudaFree(d_temp_storage));
 
-        int src = (is_reverse_edge == true) ? std::get<1>(edge) : std::get<0>(edge);
-        int dst = (is_reverse_edge == true) ? std::get<0>(edge) : std::get<1>(edge);
+    // Invoking the label edges kernel
+    SIZE_TYPE THREADS_NUM = 128;
+    SIZE_TYPE BLOCKS_NUM = CALC_BLOCKS_NUM(THREADS_NUM, (SIZE_TYPE)gpma.row_num);
+    label_edges_kernel<<<BLOCKS_NUM,THREADS_NUM>>>(RAW_PTR(gpma.row_offset), RAW_PTR(gpma.keys), RAW_PTR(gpma.values), RAW_PTR(gpma.cum_out_degree), gpma.row_num);
+    cErr(cudaDeviceSynchronize());
+}
 
-        // updating the in-degree and out-degree of
-        // destination and source node respectively
-        if (is_delete)
-        {
-            gpma.in_degree[dst] -= 1;
-            gpma.out_degree[src] -= 1;
-        }
-        else
-        {
-            gpma.in_degree[dst] += 1;
-            gpma.out_degree[src] += 1;
+__global__ void count_sort_kernel(SIZE_TYPE* bwd_row_offset, KEY_TYPE* bwd_column_indices, VALUE_TYPE* bwd_values, SIZE_TYPE* fwd_row_offset, KEY_TYPE* fwd_column_indices, VALUE_TYPE* fwd_values, int row_offset_size){
+    // [WARNING 1] HERE (src,pos) WE ARE INDEXING WITH UNISGNED INT
+    // THERE ARE POTENTIAL ISSUES THAT NEED TO BE CONSIDERED
+    // [WARNING 2] THIS IS NO LONGER A STABLE SORT
+    
+    int index = blockDim.x * blockIdx.x + threadIdx.x;
+    int block_offset = gridDim.x * blockDim.x;
+
+    for(; index < row_offset_size-1; index += block_offset){
+        int beg = fwd_row_offset[index];
+        int end = fwd_row_offset[index+1];
+        for(int i=beg; i<end; ++i){
+            KEY_TYPE key = fwd_column_indices[i];
+            VALUE_TYPE value = fwd_values[i];
+            SIZE_TYPE src = key;
+            if (key != KEY_MAX && src != COL_IDX_NONE && value != VALUE_NONE)
+            {
+                SIZE_TYPE pos = atomicSub(&bwd_row_offset[src],1) - 1;
+                bwd_column_indices[pos] = ((KEY_TYPE)src << 32) + (key >> 32);
+                bwd_values[pos] = value;
+            }
         }
+    }
+}
 
-        host_src.push_back(src);
-        host_dst.push_back(dst);
+std::vector<float> build_backward_csr(GPMA &gpma){
 
-        ++edge_count;
-    }
+    // Step 1: Allocate space for backward CSR vectors
+    auto start_time_alloc = std::chrono::high_resolution_clock::now();
+    cErr(cudaMalloc(&gpma.bwd_row_offset, (gpma.row_num+1) * sizeof(SIZE_TYPE)));
+    cErr(cudaMalloc(&gpma.bwd_keys, gpma.edge_count * sizeof(KEY_TYPE)));
+    cErr(cudaMalloc(&gpma.bwd_values, gpma.edge_count * sizeof(VALUE_TYPE)));
+    auto end_time_alloc = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<float> time_alloc = (end_time_alloc - start_time_alloc);
+
+    // Step 2: Find the cummulative sum of the in-degree vector
+    auto start_time_cum = std::chrono::high_resolution_clock::now();
+    void *d_temp_storage = NULL;
+    size_t temp_storage_bytes = 0;
+    cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, RAW_PTR(gpma.in_degree), gpma.bwd_row_offset, gpma.row_num);
+    cudaMalloc(&d_temp_storage, temp_storage_bytes);
+    cub::DeviceScan::InclusiveSum(d_temp_storage, temp_storage_bytes, RAW_PTR(gpma.in_degree), gpma.bwd_row_offset, gpma.row_num);
+    cErr(cudaDeviceSynchronize());
+    cErr(cudaFree(d_temp_storage));
 
-    gpma.edge_count = is_delete ? gpma.edge_count - edge_count : gpma.edge_count + edge_count;
+    // Step 3: Insert the edge count in the last element slot in bwd_row_offset
+    SIZE_TYPE edge_count = gpma.edge_count;
+    cudaMemcpy(gpma.bwd_row_offset + gpma.row_num, &edge_count, sizeof(SIZE_TYPE), cudaMemcpyHostToDevice);
+    auto end_time_cum = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<float> time_cum = (end_time_cum - start_time_cum);
 
-    thrust::host_vector<KEY_TYPE> h_base_keys(edge_count);
+    // Step 4: Invoke the count_sort_kernel to perform the sort using the
+    // last 32-bits as key
+    auto start_time_count = std::chrono::high_resolution_clock::now();
+    SIZE_TYPE THREADS_NUM = 128;
+    SIZE_TYPE BLOCKS_NUM = CALC_BLOCKS_NUM(THREADS_NUM, (SIZE_TYPE)gpma.keys.size());
+    count_sort_kernel<<<BLOCKS_NUM,THREADS_NUM>>>(gpma.bwd_row_offset, gpma.bwd_keys, gpma.bwd_values, RAW_PTR(gpma.row_offset), RAW_PTR(gpma.keys), RAW_PTR(gpma.values), gpma.row_offset.size());
+    cErr(cudaDeviceSynchronize());
+    auto end_time_count = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<float> time_count = (end_time_count - start_time_count);
+
+    std::vector<float> res;
+    res.push_back(time_alloc.count());
+    res.push_back(time_cum.count());
+    res.push_back(time_count.count());
+    return res;
+}
 
-    for (int i = 0; i < edge_count; i++)
-        h_base_keys[i] = ((KEY_TYPE)host_src[i] << 32) + host_dst[i];
+void free_backward_csr(GPMA &gpma){
+    cErr(cudaFree(gpma.bwd_row_offset));
+    cErr(cudaFree(gpma.bwd_keys));
+    cErr(cudaFree(gpma.bwd_values));
+}
 
-    DEV_VEC_KEY base_keys = h_base_keys;
-    cudaDeviceSynchronize();
+std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> get_csr_ptrs(GPMA &gpma, bool is_backward = false)
+{
+    // This function returns CSR pointers to forward or backward graph
+    // based on if is_backward is false or true respectively.
 
-    DEV_VEC_VALUE base_values(edge_count);
+    std::tuple<std::size_t, std::size_t, std::size_t> t;
+    if(is_backward){
+        std::get<0>(t) = (std::uintptr_t)gpma.bwd_row_offset;
+        std::get<1>(t) = (std::uintptr_t)gpma.bwd_keys;
+        std::get<2>(t) = (std::uintptr_t)gpma.bwd_values;
+    }else{
+        std::get<0>(t) = (std::uintptr_t)RAW_PTR(gpma.row_offset);
+        std::get<1>(t) = (std::uintptr_t)RAW_PTR(gpma.keys);
+        std::get<2>(t) = (std::uintptr_t)RAW_PTR(gpma.values);
+    }
+    return t;
+}
 
-    if (is_delete)
-        thrust::fill(base_values.begin(), base_values.end(), VALUE_NONE);
-    else
-        thrust::fill(base_values.begin(), base_values.end(), 1);
+std::vector<SIZE_TYPE> get_in_degrees(GPMA &gpma){
+    cErr(cudaMemcpy(gpma.in_degree_host, RAW_PTR(gpma.in_degree), sizeof(SIZE_TYPE) * gpma.row_num, cudaMemcpyDeviceToHost));
+    std::vector<SIZE_TYPE> in_degree(gpma.in_degree_host, gpma.in_degree_host + gpma.row_num);
+    return in_degree;
+} 
 
-    cudaDeviceSynchronize();
+std::vector<SIZE_TYPE> get_out_degrees(GPMA &gpma){
+    cErr(cudaMemcpy(gpma.out_degree_host, RAW_PTR(gpma.out_degree), sizeof(SIZE_TYPE) * gpma.row_num, cudaMemcpyDeviceToHost));
+    std::vector<SIZE_TYPE> out_degree(gpma.out_degree_host, gpma.out_degree_host + gpma.row_num);
+    return out_degree;
+} 
 
-    update_gpma(gpma, base_keys, base_values);
-    cudaDeviceSynchronize();
+//////////////////////////////////////// GPMA Logging APIs ///////////////////////////////////////////////////////////
+
+std::tuple<int, int> get_graph_attr(GPMA &gpma)
+{
+    std::tuple<int, int> t;
+
+    // Since row_offset has size num_nodes + 1
+    std::get<0>(t) = gpma.row_offset.size() - 1;
+    std::get<1>(t) = gpma.edge_count;
+    return t;
 }
 
-void label_edges(GPMA &gpma)
+void print_gpma_info(GPMA &gpma, int node)
 {
-    int edge_label_counter = 1;
+    // Pretty prints the row_offset range and column_indices values
+    // for the given node
+    //
+    // row_offset:  Row offset vector
+    // col_indices: Column indices PMA
+    // node:        Node whose row_offset and column_indices is to be printed
+
+    // getting the (beg,end) range of the node
 
-    thrust::host_vector<SIZE_TYPE> h_row_offset = gpma.row_offset;
-    thrust::host_vector<KEY_TYPE> h_keys = gpma.keys;
-    thrust::host_vector<VALUE_TYPE> h_values = gpma.values;
+    thrust::host_vector<SIZE_TYPE> row_offset = gpma.row_offset;
+    thrust::host_vector<KEY_TYPE> col_indices = gpma.keys;
+    thrust::host_vector<VALUE_TYPE> edge_values = gpma.values;
+
+    unsigned int beg = row_offset[node];
+    unsigned int end = row_offset[node + 1];
+    int print_count = 0;
+
+    py::print("\n🌟 Column Indices for node: ", node, "(", beg, ", ", end, ")", "\n\n");
+    // std::cout << std::setw(6);
 
-    for (int node = 0; node < h_row_offset.size() - 1; ++node)
+    for (int i = beg; i < end; ++i)
     {
-        unsigned int beg = h_row_offset[node];
-        unsigned int end = h_row_offset[node + 1];
+        // KEY_TYPE mask = (KEY_TYPE)node << 32;
+        // unsigned int dst = (col_indices[i] - mask);
+        unsigned int dst = (col_indices[i] & 0xffffffff);
+        VALUE_TYPE val = edge_values[i];
 
-        for (int i = beg; i < end; ++i)
+        if (col_indices[i] != KEY_MAX && (col_indices[i] & COL_IDX_NONE != COL_IDX_NONE) && val != VALUE_NONE)
         {
-            KEY_TYPE mask = (KEY_TYPE)node << 32;
-            unsigned int dst = (h_keys[i] - mask);
-            if (dst != COL_IDX_NONE && h_values[i] != VALUE_NONE)
-            {
-                h_values[i] = edge_label_counter;
-                edge_label_counter += 1;
-            }
+            py::print("[", col_indices[i], "] ", dst, "(", val, ")", "  ");
         }
-    }
+        else
+        {
+            py::print("[", col_indices[i], "] ", "(-)(", val, ")", "  ");
+        }
+        print_count += 1;
 
-    gpma.row_offset = h_row_offset;
-    gpma.keys = h_keys;
-    gpma.values = h_values;
+        if (print_count > 10)
+        {
+            py::print("\n");
+            print_count = 0;
+        }
+    }
+    py::print("\n\n");
 }
 
-void copy_label_edges(GPMA &gpma, GPMA &ref_gpma)
+std::set<std::tuple<unsigned int, unsigned int, unsigned int>> get_gpma_edge_list(GPMA &gpma)
 {
-    int edge_counter = 0;
 
-    // creating the host vector version for keys
-    // and values for the new gpma
-    thrust::host_vector<KEY_TYPE> h_new_keys(ref_gpma.edge_count);
-    thrust::host_vector<VALUE_TYPE> h_new_values(ref_gpma.edge_count);
+    std::set<std::tuple<unsigned int, unsigned int, unsigned int>> vec;
+    thrust::host_vector<SIZE_TYPE> h_ref_row_offset = gpma.row_offset;
+    thrust::host_vector<KEY_TYPE>
+        h_ref_keys = gpma.keys;
+    thrust::host_vector<VALUE_TYPE> h_ref_values = gpma.values;
 
-    // creating the host vectors for the CSR arrays of the
-    // reference GPMA from it's device vectors
-    thrust::host_vector<SIZE_TYPE> h_ref_row_offset = ref_gpma.row_offset;
-    thrust::host_vector<KEY_TYPE> h_ref_keys = ref_gpma.keys;
-    thrust::host_vector<VALUE_TYPE> h_ref_values = ref_gpma.values;
+    cErr(cudaDeviceSynchronize());
 
-    DEV_VEC_KEY d_new_keys(ref_gpma.edge_count);
-    DEV_VEC_VALUE d_new_values(ref_gpma.edge_count);
+    // py::print("ROW_OFFSET: ", h_ref_row_offset.size());
+    // py::print("COL_IDX_SIZE: ", h_ref_keys.size());
+    // py::print("VALUES_SIZE: ", h_ref_values.size());
 
     for (int node = 0; node < h_ref_row_offset.size() - 1; ++node)
     {
-        unsigned int beg = h_ref_row_offset[node];
-        unsigned int end = h_ref_row_offset[node + 1];
-
-        for (int i = beg; i < end; ++i)
+        SIZE_TYPE beg = h_ref_row_offset[node];
+        SIZE_TYPE end = h_ref_row_offset[node + 1];
+        for (SIZE_TYPE i = beg; i < end; ++i)
         {
-            KEY_TYPE mask = (KEY_TYPE)node << 32;
-            unsigned int dst = (h_ref_keys[i] - mask);
-            if (dst != COL_IDX_NONE && h_ref_values[i] != VALUE_NONE)
+            // h_ref_keys[i] != KEY_MAX && h_ref_keys[i] != KEY_NONE
+            // KEY_TYPE mask = (KEY_TYPE)node << 32;
+            // unsigned int dst = (h_ref_keys[i] - mask);
+            unsigned int dst = (h_ref_keys[i] & 0xffffffff);
+            if (h_ref_keys[i] != KEY_MAX && dst != COL_IDX_NONE && h_ref_values[i] != VALUE_NONE)
             {
-                h_new_keys[edge_counter] = (h_ref_keys[i] << 32) + node;
-                h_new_values[edge_counter] = h_ref_values[i];
-                edge_counter += 1;
+                std::tuple<unsigned int, unsigned int, unsigned int> tup;
+                unsigned int src = (h_ref_keys[i] >> 32);
+                std::get<0>(tup) = src;
+                std::get<1>(tup) = dst;
+                std::get<2>(tup) = h_ref_values[i];
+                vec.insert(tup);
             }
         }
     }
 
-    d_new_keys = h_new_keys;
-    d_new_values = h_new_values;
-
-    // NOTE: Verify if sorting is really required or not
-    thrust::sort_by_key(d_new_keys.begin(), d_new_keys.end(), d_new_values.begin());
-    cudaDeviceSynchronize();
-
-    locate_leaf_batch(RAW_PTR(gpma.keys), RAW_PTR(gpma.values), gpma.keys.size(), gpma.segment_length, gpma.tree_height,
-                      RAW_PTR(d_new_keys), RAW_PTR(d_new_values), d_new_keys.size(), NULL, false);
-
-    cudaDeviceSynchronize();
+    return vec;
 }
 
-void build_reverse_gpma(GPMA &gpma, GPMA &ref_gpma)
+std::set<std::tuple<unsigned int, unsigned int, unsigned int>> get_reverse_csr_edge_list(GPMA &gpma)
 {
-    // This funtion expects an empty initialized GPMA
-    // as first param and the reference GPMA from which
-    // edges are to be added as the second param
+    SIZE_TYPE* h_ref_row_offset = (SIZE_TYPE*) malloc((gpma.row_num+1) * sizeof(SIZE_TYPE));
+    KEY_TYPE* h_ref_keys = (KEY_TYPE*) malloc((gpma.edge_count) * sizeof(KEY_TYPE));
+    VALUE_TYPE* h_ref_values = (VALUE_TYPE*) malloc((gpma.edge_count) * sizeof(VALUE_TYPE));
 
-    int edge_counter = 0;
 
-    // creating host_vectors for the new GPMA
-    thrust::host_vector<KEY_TYPE> h_new_keys(ref_gpma.edge_count);
-    thrust::host_vector<VALUE_TYPE> h_new_values(ref_gpma.edge_count);
+    cErr(cudaMemcpy(h_ref_row_offset, gpma.bwd_row_offset, (gpma.row_num+1) * sizeof(SIZE_TYPE), cudaMemcpyDeviceToHost));
+    cErr(cudaMemcpy(h_ref_keys, gpma.bwd_keys, (gpma.edge_count) * sizeof(KEY_TYPE), cudaMemcpyDeviceToHost));
+    cErr(cudaMemcpy(h_ref_values, gpma.bwd_values, (gpma.edge_count) * sizeof(VALUE_TYPE), cudaMemcpyDeviceToHost));
 
-    // creating the host vectors for the CSR arrays of the
-    // reference GPMA from it's device vectors
-    thrust::host_vector<SIZE_TYPE> h_ref_row_offset = ref_gpma.row_offset;
-    thrust::host_vector<KEY_TYPE> h_ref_keys = ref_gpma.keys;
-    thrust::host_vector<VALUE_TYPE> h_ref_values = ref_gpma.values;
+    std::set<std::tuple<unsigned int, unsigned int, unsigned int>> vec;
 
-    DEV_VEC_KEY d_new_keys(ref_gpma.edge_count);
-    DEV_VEC_VALUE d_new_values(ref_gpma.edge_count);
+    // py::print("ROW_OFFSET: ", h_ref_row_offset.size());
+    // py::print("COL_IDX_SIZE: ", h_ref_keys.size());
+    // py::print("VALUES_SIZE: ", h_ref_values.size());
 
-    for (int node = 0; node < h_ref_row_offset.size() - 1; ++node)
+    for (int node = 0; node < gpma.row_num; ++node)
     {
-        unsigned int beg = h_ref_row_offset[node];
-        unsigned int end = h_ref_row_offset[node + 1];
-
-        for (int i = beg; i < end; ++i)
+        SIZE_TYPE beg = h_ref_row_offset[node];
+        SIZE_TYPE end = h_ref_row_offset[node + 1];
+        for (SIZE_TYPE i = beg; i < end; ++i)
         {
-            KEY_TYPE mask = (KEY_TYPE)node << 32;
-            unsigned int dst = (h_ref_keys[i] - mask);
-            if (dst != COL_IDX_NONE && h_ref_values[i] != VALUE_NONE)
+            // h_ref_keys[i] != KEY_MAX && h_ref_keys[i] != KEY_NONE
+            // KEY_TYPE mask = (KEY_TYPE)node << 32;
+            // unsigned int dst = (h_ref_keys[i] - mask);
+            unsigned int dst = (h_ref_keys[i] & 0xffffffff);
+            if (h_ref_keys[i] != KEY_MAX && dst != COL_IDX_NONE && h_ref_values[i] != VALUE_NONE)
             {
-                h_new_keys[edge_counter] = (h_ref_keys[i] << 32) + node;
-                h_new_values[edge_counter] = h_ref_values[i];
-                edge_counter += 1;
+                std::tuple<unsigned int, unsigned int, unsigned int> tup;
+                unsigned int src = (h_ref_keys[i] >> 32);
+                std::get<0>(tup) = src;
+                std::get<1>(tup) = dst;
+                std::get<2>(tup) = h_ref_values[i];
+                vec.insert(tup);
             }
         }
     }
 
-    d_new_keys = h_new_keys;
-    d_new_values = h_new_values;
-
-    update_gpma(gpma, d_new_keys, d_new_values);
-    cudaDeviceSynchronize();
-}
-
-std::tuple<std::size_t, std::size_t, std::size_t> get_csr_ptrs(GPMA &gpma)
-{
-    std::tuple<std::size_t, std::size_t, std::size_t> t;
-    std::get<0>(t) = (std::size_t)RAW_PTR(gpma.row_offset);
-    std::get<1>(t) = (std::size_t)RAW_PTR(gpma.keys);
-    std::get<2>(t) = (std::size_t)RAW_PTR(gpma.values);
-    return t;
-}
-
-std::tuple<int, int> get_graph_attr(GPMA &gpma)
-{
-    std::tuple<int, int> t;
-
-    // Since row_offset has size num_nodes + 1
-    std::get<0>(t) = gpma.row_offset.size() - 1;
-    std::get<1>(t) = gpma.edge_count;
-    return t;
+    return vec;
 }
 
 PYBIND11_MODULE(gpma, m)
 {
     m.doc() = "CPython module for GPMA"; // optional module docstring
 
-    m.def("init_gpma", &init_gpma, "Initialises the CSR arrays using GPMA");
-    m.def("print_gpma_info", &print_gpma_info, "Prints row_offset and col_indices for a given node");
-    m.def("load_graph", &load_graph, "Loads a graph data into a GPMA");
-    m.def("edge_update_list", &edge_update_list, "Updates the GPMA by adding/deleting edges from the edge list", py::arg("gpma"), py::arg("edge_list"), py::arg("is_delete") = false, py::arg("is_reverse_edge") = false);
-    m.def("label_edges", &label_edges, "Creates edge labels for the current GPMA");
-    m.def("copy_label_edges", &copy_label_edges, "Label edges of a GPMA based on another GPMA");
-    m.def("build_reverse_gpma", &build_reverse_gpma, "Builds the reverse GPMA based on another GPMA");
-    m.def("get_csr_ptrs", &get_csr_ptrs, "Returns the pointers to row_offset, col_indices and edge_ids");
+    // GPMA APIs
+    m.def("init_gpma", &init_gpma, "Initialises the CSR arrays using GPMA", py::arg("gpma"), py::arg("num_nodes"));
+    m.def("init_graph_updates", &init_graph_updates, "Initialize graph updates", py::arg("gpma"), py::arg("updates"), py::arg("reverse_edges") = false);
+    m.def("edge_update_t", &edge_update_t, "Performs the updates of the given timestamp", py::arg("gpma"), py::arg("timestamp"), py::arg("revert_update") = false);
+    m.def("label_edges", &label_edges, "Labels edges of the given GPMA", py::arg("gpma"));
+    m.def("build_backward_csr", &build_backward_csr, "Build the backward CSR for given GPMA", py::arg("gpma"));
+    m.def("free_backward_csr", &free_backward_csr, "Free backward CSR resources from given GPMA", py::arg("gpma"));
+    m.def("get_csr_ptrs", &get_csr_ptrs, "Returns the pointers to row_offset, col_indices and edge_ids", py::arg("gpma"), py::arg("is_backward") = false);
+    m.def("get_in_degrees", &get_in_degrees, "Returns the in-degrees of all nodes in given GPMA", py::arg("gpma"));
+    m.def("get_out_degrees", &get_out_degrees, "Returns the out-degrees of all nodes in given GPMA", py::arg("gpma"));
+    
+    // GPMA Logging APIs
     m.def("get_graph_attr", &get_graph_attr, "Returns the [num_nodes, num_edges]");
+    m.def("print_gpma_info", &print_gpma_info, "Prints row_offset and col_indices for a given node");
+    m.def("get_gpma_edge_list", &get_gpma_edge_list, "To get the edge list");
+    m.def("get_reverse_csr_edge_list", &get_reverse_csr_edge_list, "To get the edge list of reverse");
 
     py::class_<GPMA>(m, "GPMA")
         .def(py::init<>())
-        .def("get_size", &GPMA::get_size)
-        .def_readwrite("keys", &GPMA::keys)
-        .def_readwrite("values", &GPMA::values)
-        .def_readwrite("segment_length", &GPMA::segment_length)
-        .def_readwrite("tree_height", &GPMA::tree_height)
-        .def_readwrite("density_lower_thres_leaf", &GPMA::density_lower_thres_leaf)
-        .def_readwrite("density_lower_thres_root", &GPMA::density_lower_thres_root)
-        .def_readwrite("density_upper_thres_root", &GPMA::density_upper_thres_root)
-        .def_readwrite("density_upper_thres_leaf", &GPMA::density_upper_thres_leaf)
-        .def_readwrite("lower_element", &GPMA::lower_element)
-        .def_readwrite("upper_element", &GPMA::upper_element)
-        .def_readwrite("row_num", &GPMA::row_num)
-        .def_readwrite("row_offset", &GPMA::row_offset)
-        .def_readwrite("in_degree", &GPMA::in_degree)
-        .def_readwrite("out_degree", &GPMA::out_degree)
         .def("__copy__", [](const GPMA &self)
              { return GPMA(self); })
         .def(
@@ -1340,5 +1434,5 @@ PYBIND11_MODULE(gpma, m)
             "memo"_a);
 }
 
-// Command used:
+// Command used for compilation:
 // /usr/local/cuda-11.7/bin/nvcc $(python3 -m pybind11 --includes) -shared -rdc=true --compiler-options '-fPIC'  -D__CDPRT_SUPPRESS_SYNC_DEPRECATION_WARNING -o gpma.so gpma.cu
\ No newline at end of file
diff --git a/seastar/graph/dynamic/naive/NaiveGraph.py b/seastar/graph/dynamic/naive/NaiveGraph.py
index f87bf9b..4b53f18 100644
--- a/seastar/graph/dynamic/naive/NaiveGraph.py
+++ b/seastar/graph/dynamic/naive/NaiveGraph.py
@@ -7,83 +7,95 @@
 from collections import deque
 import time
 
+
 class NaiveGraph(DynamicGraph):
-    def __init__(self, edge_list):
-        super().__init__(edge_list)
+    def __init__(self, edge_list, max_num_nodes):
+        super().__init__(edge_list, max_num_nodes)
         # inspect(edge_list)
         self._prepare_edge_lst_fwd(edge_list)
-        self._prepare_edge_lst_bwd(self.fwd_edge_list)  
-        self._forward_graph = [CSR(self.fwd_edge_list[i], self.graph_updates[str(i)]["num_nodes"], is_edge_reverse=True) for i in range(len(self.fwd_edge_list))]
-        self._backward_graph = [CSR(self.bwd_edge_list[i], self.graph_updates[str(i)]["num_nodes"]) for i in range(len(self.bwd_edge_list))]
-        # self._forward_graph = [CSR(self.fwd_edge_list[0], self.graph_updates[str(0)]["num_nodes"], is_edge_reverse=True)]
-        # self._backward_graph = [CSR(self.bwd_edge_list[0], self.graph_updates[str(0)]["num_nodes"])]
+        self._prepare_edge_lst_bwd(self.fwd_edge_list)
+
+        # This is used because edge weights is a compulsary argument to CSR
+        edge_weight_lst = [[1 for _ in edge_list_t] for edge_list_t in edge_list]
+
+        self._forward_graph = [
+            CSR(
+                self.fwd_edge_list[i],
+                edge_weight_lst[i],
+                self.graph_attr[str(i)][0],
+                is_edge_reverse=True,
+            )
+            for i in range(len(self.fwd_edge_list))
+        ]
+        self._backward_graph = [
+            CSR(self.bwd_edge_list[i], edge_weight_lst[i], self.graph_attr[str(i)][0])
+            for i in range(len(self.bwd_edge_list))
+        ]
+
+        # for benchmarking purposes
+        self._update_count = 0
+        self._total_update_time = 0
+        self._gpu_move_time = 0
 
         self._get_graph_csr_ptrs(0)
-        
-    def _prepare_edge_lst_fwd(self, edge_list):   
+
+    def _prepare_edge_lst_fwd(self, edge_list):
         self.fwd_edge_list = []
         for i in range(len(edge_list)):
             edge_list_for_t = edge_list[i]
-            edge_list_for_t.sort(key = lambda x: (x[1],x[0]))
-            edge_list_for_t = [(edge_list_for_t[j][0],edge_list_for_t[j][1],j) for j in range(len(edge_list_for_t))]
+            edge_list_for_t.sort(key=lambda x: (x[1], x[0]))
+            edge_list_for_t = [
+                (edge_list_for_t[j][0], edge_list_for_t[j][1], j)
+                for j in range(len(edge_list_for_t))
+            ]
             self.fwd_edge_list.append(edge_list_for_t)
-    
-    def _prepare_edge_lst_bwd(self, edge_list):    
+
+    def _prepare_edge_lst_bwd(self, edge_list):
         self.bwd_edge_list = []
         for i in range(len(edge_list)):
             edge_list_for_t = copy.deepcopy(edge_list[i])
             edge_list_for_t.sort()
             self.bwd_edge_list.append(edge_list_for_t)
-        
-    # def _graph_stack_push(self, elem):
-    #     self.graph_stack.append(elem)
-    
-    # def _graph_stack_pop(self):
-    #     self.graph_stack.pop()
-        
-    # def _graph_stack_top(self):
-    #     return self.graph_stack[-1]
-        
+
     def graph_type(self):
         return "csr"
-        
+
     def in_degrees(self):
-        return np.array(self._forward_graph[self.current_timestamp].out_degrees, dtype='int32')
+        return np.array(self._forward_graph[self.current_timestamp].out_degrees, dtype="int32")
     
     def out_degrees(self):
-        return np.array(self._forward_graph[self.current_timestamp].in_degrees, dtype='int32')
+        return np.array(self._forward_graph[self.current_timestamp].in_degrees, dtype="int32")
     
+    # def weighted_in_degrees(self):
+    #     return np.array(self._forward_graph[self.current_timestamp].weighted_out_degrees, dtype="float32")
+
     def _get_graph_csr_ptrs(self, timestamp):
         if self._is_backprop_state:
             bwd_csr_ptrs = self._backward_graph[timestamp]
-            # bwd_csr_ptrs.print_csr_arrays()
             self.bwd_row_offset_ptr = bwd_csr_ptrs.row_offset_ptr
             self.bwd_column_indices_ptr = bwd_csr_ptrs.column_indices_ptr
             self.bwd_eids_ptr = bwd_csr_ptrs.eids_ptr
         else:
             fwd_csr_ptrs = self._forward_graph[timestamp]
-            # fwd_csr_ptrs.print_csr_arrays()
             self.fwd_row_offset_ptr = fwd_csr_ptrs.row_offset_ptr
             self.fwd_column_indices_ptr = fwd_csr_ptrs.column_indices_ptr
             self.fwd_eids_ptr = fwd_csr_ptrs.eids_ptr
 
-            # print(f"(PYTHON) RECEIVED ROW OFFSET PTR: {fwd_csr_ptrs.row_offset_ptr}")
-            # print_dev_array(fwd_csr_ptrs.row_offset_ptr,129)
-            # print("QUITTING")
-            # quit()
-    
     def _update_graph_forward(self):
-        ''' Updates the current base graph to the next timestamp
-        '''
+        """Updates the current base graph to the next timestamp"""
         if str(self.current_timestamp + 1) not in self.graph_updates:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_forward()")
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_forward()"
+            )
         self._get_graph_csr_ptrs(self.current_timestamp + 1)
-        
+
     def _init_reverse_graph(self):
-        ''' Generates the reverse of the base graph'''
+        """Generates the reverse of the base graph"""
         self._get_graph_csr_ptrs(self.current_timestamp)
-        
+
     def _update_graph_backward(self):
         if self.current_timestamp < 0:
-            raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_backward()")
+            raise Exception(
+                "⏰ Invalid timestamp during SeastarGraph.update_graph_backward()"
+            )
         self._get_graph_csr_ptrs(self.current_timestamp - 1)
diff --git a/seastar/graph/dynamic/pcsr/PCSRGraph.py b/seastar/graph/dynamic/pcsr/PCSRGraph.py
index 3263835..0b81fc9 100644
--- a/seastar/graph/dynamic/pcsr/PCSRGraph.py
+++ b/seastar/graph/dynamic/pcsr/PCSRGraph.py
@@ -4,82 +4,48 @@
 import time
 
 from seastar.graph.dynamic.DynamicGraph import DynamicGraph
-from seastar.graph.dynamic.pcsr.pcsr import PCSR, build_reverse_pcsr
+from seastar.graph.dynamic.pcsr.pcsr import PCSR
 
 class PCSRGraph(DynamicGraph):
-    def __init__(self, edge_list):
-        super().__init__(edge_list)
+    def __init__(self, edge_list, max_num_nodes):
+        super().__init__(edge_list, max_num_nodes)
         
-        # preprocessing the edge ids for forward and backward pass
-        # so that it doesn't have to be calculated each time
-        self._prepare_eid_bwd(edge_list) 
+        # Get the maximum number of edges
+        self._get_max_num_edges()
         
-        self._forward_graph = PCSR(self.max_num_nodes)
-        self._backward_graph = PCSR(self.max_num_nodes)
-        self._forward_graph.edge_update_list(self.graph_updates["0"]["add"],is_reverse_edge=True)
-        
-        # Using cache to possibly make pcsr faster
-        self._in_degrees_cache = {}
-        self._out_degrees_cache = {}
-        
-        self._update_graph_cache()
-        self._get_graph_csr_ptrs(eids=list())
+        self._forward_graph = PCSR(self.max_num_nodes, self.max_num_edges)
+        self._forward_graph.edge_update_list(self.graph_updates["0"]["add"], is_reverse_edge=True)
+        self._forward_graph.label_edges()   
+        self._forward_graph.build_csr()
+        self._get_graph_csr_ptrs()
     
-    def _prepare_eid_bwd(self, edge_list):   
-        fwd_edge_list = []
-        for i in range(len(edge_list)):
-            edge_list_for_t = edge_list[i]
-            edge_list_for_t.sort(key = lambda x: (x[1],x[0]))
-            edge_list_for_t = [(edge_list_for_t[j][0],edge_list_for_t[j][1],j) for j in range(len(edge_list_for_t))]
-            fwd_edge_list.append(edge_list_for_t)
-        
-        self.bwd_eid_list = []
-        for i in range(len(fwd_edge_list)):
-            edge_list_for_t = fwd_edge_list[i]
-            edge_list_for_t.sort()
-            eid_t = [edge_list_for_t[i][2] for i in range(len(edge_list_for_t))]
-            self.bwd_eid_list.append(eid_t)
+    def _get_max_num_edges(self):
+        updates = self.graph_updates
+        edge_set = set()
+        for i in range(len(updates)):
+          for j in range(len(updates[str(i)]["add"])):
+            edge_set.add(updates[str(i)]["add"][j])
+        self.max_num_edges = len(edge_set)
                 
     def graph_type(self):
         return "pcsr"
         
     def in_degrees(self):
-        if self.current_timestamp not in self._in_degrees_cache:
-            self._in_degrees_cache[self.current_timestamp] = np.array([node.num_neighbors for node in self._forward_graph.nodes], dtype='int32')
-        
-        return self._in_degrees_cache[self.current_timestamp]
+        return np.array(self._forward_graph.out_degrees, dtype='int32')
     
     def out_degrees(self):
-        if self.current_timestamp not in self._out_degrees_cache:
-            self._out_degrees_cache[self.current_timestamp] = np.array([node.in_degree for node in self._forward_graph.nodes], dtype='int32')
-        
-        return self._out_degrees_cache[self.current_timestamp]
+        return np.array(self._forward_graph.in_degrees, dtype='int32')
     
-    def _update_graph_cache(self, is_bwd=False):
-        if is_bwd:
-            # saving reverse base graph in cache
-            self.graph_cache['bwd'] = copy.deepcopy(self._backward_graph)
-        else:
-            self.graph_cache['fwd'] = copy.deepcopy(self._forward_graph)
-            
-    def _get_cached_graph(self, is_bwd=False):
-        if is_bwd:
-            return copy.deepcopy(self.graph_cache['bwd'])
-        else:
-            return copy.deepcopy(self.graph_cache['fwd'])
-    
-    def _get_graph_csr_ptrs(self, eids): 
-
+    def _get_graph_csr_ptrs(self): 
+        csr_ptrs = self._forward_graph.get_csr_ptrs()
         if self._is_backprop_state:
-            backward_csr_ptrs = self._backward_graph.get_csr_ptrs(eids=eids)
-            self.bwd_row_offset_ptr = backward_csr_ptrs[0]
-            self.bwd_column_indices_ptr = backward_csr_ptrs[1]
-            self.bwd_eids_ptr = backward_csr_ptrs[2]
+            self.bwd_row_offset_ptr = csr_ptrs[0]
+            self.bwd_column_indices_ptr = csr_ptrs[1]
+            self.bwd_eids_ptr = csr_ptrs[2]
         else:
-            forward_csr_ptrs = self._forward_graph.get_csr_ptrs(eids=eids)
-            self.fwd_row_offset_ptr = forward_csr_ptrs[0]
-            self.fwd_column_indices_ptr = forward_csr_ptrs[1]
-            self.fwd_eids_ptr = forward_csr_ptrs[2]
+            self.fwd_row_offset_ptr = csr_ptrs[0]
+            self.fwd_column_indices_ptr = csr_ptrs[1]
+            self.fwd_eids_ptr = csr_ptrs[2]
     
     def _update_graph_forward(self):
         ''' Updates the current base graph to the next timestamp
@@ -89,43 +55,27 @@ def _update_graph_forward(self):
         
         graph_additions = self.graph_updates[str(self.current_timestamp + 1)]["add"]
         graph_deletions = self.graph_updates[str(self.current_timestamp + 1)]["delete"]
-        
-        # print(f"(FORWARD) ➕➕➕ Adding a total of {len(graph_additions)} elements", flush=True)
-        # print(f"(FORWARD) ➖➖➖ Removing a total of {len(graph_deletions)} elements", flush=True)
 
-        self._forward_graph.edge_update_list(graph_additions,is_reverse_edge=True)
-        self._forward_graph.edge_update_list(graph_deletions,is_delete=True,is_reverse_edge=True)
-        self._get_graph_csr_ptrs(eids=list())
+        self._forward_graph.edge_update_list(graph_additions, is_reverse_edge=True)
+        self._forward_graph.edge_update_list(graph_deletions, is_delete=True, is_reverse_edge=True)
+        self._forward_graph.label_edges()   
+        self._forward_graph.build_csr()
+        self._get_graph_csr_ptrs()
         
     def _init_reverse_graph(self):
         ''' Generates the reverse of the base graph'''
-        
-        # checking if the reverse base graph exists in the cache
-        # we can load it from there instead of building it each time
-        if 'bwd' in self.graph_cache:
-            self._backward_graph = self._get_cached_graph(is_bwd=True)
-        else:
-            build_reverse_pcsr(self._backward_graph, self._forward_graph)
+        self._forward_graph.build_reverse_csr()
+        self._get_graph_csr_ptrs()
 
-            # storing the reverse base graph in cache after building
-            # it for the first time
-            self._update_graph_cache(is_bwd=True)
-            
-        self._get_graph_csr_ptrs(eids=self.bwd_eid_list[self.current_timestamp])
-        
-        # Resetting the forward graph so that it can be used later
-        self._forward_graph = self._get_cached_graph()
-        
     def _update_graph_backward(self):
         if self.current_timestamp < 0:
             raise Exception("⏰ Invalid timestamp during SeastarGraph.update_graph_backward()")
         
         graph_additions = self.graph_updates[str(self.current_timestamp)]["delete"]
         graph_deletions = self.graph_updates[str(self.current_timestamp)]["add"]
-        
-        # print(f"(BACKWARD) ➕➕➕ Adding a total of {len(graph_additions)} elements",flush=True)
-        # print(f"(BACKWARD) ➖➖➖ Removing a total of {len(graph_deletions)} elements",flush=True)
-        
-        self._backward_graph.edge_update_list(graph_additions)   
-        self._backward_graph.edge_update_list(graph_deletions,is_delete=True)
-        self._get_graph_csr_ptrs(eids=self.bwd_eid_list[self.current_timestamp - 1])
+
+        self._forward_graph.edge_update_list(graph_additions, is_reverse_edge=True)   
+        self._forward_graph.edge_update_list(graph_deletions, is_delete=True, is_reverse_edge=True)
+        self._forward_graph.label_edges() 
+        self._forward_graph.build_reverse_csr()
+        self._get_graph_csr_ptrs()
\ No newline at end of file
diff --git a/seastar/graph/dynamic/pcsr/pcsr.cu b/seastar/graph/dynamic/pcsr/pcsr.cu
index fd1bfb1..371a9f7 100644
--- a/seastar/graph/dynamic/pcsr/pcsr.cu
+++ b/seastar/graph/dynamic/pcsr/pcsr.cu
@@ -12,18 +12,32 @@
 #include <cstring>
 #include <iostream>
 
+#include <chrono>
+
 #include <thrust/device_vector.h>
 #include <thrust/host_vector.h>
 #include <thrust/remove.h>
 #include <thrust/sort.h>
 #include <cub/cub.cuh>
+#include <thrust/system/cuda/experimental/pinned_allocator.h>
 
 namespace py = pybind11;
 using namespace pybind11::literals;
 using namespace std;
 
-typedef thrust::device_vector<int> DEV_VEC;
-#define RAW_PTR(x) thrust::raw_pointer_cast((x).data())
+using namespace std::chrono;
+
+#define cErr(errcode)                             \
+    {                                             \
+        gpuAssert((errcode), __FILE__, __LINE__); \
+    }
+__inline__ __host__ __device__ void gpuAssert(cudaError_t code, const char *file, int line)
+{
+    if (code != cudaSuccess)
+    {
+        printf("GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
+    }
+}
 
 ////////////////////////////////////////////////////////////////////////////
 
@@ -131,6 +145,7 @@ bool is_sentinel(edge_t e)
     return e.dest == UINT32_MAX || e.value == UINT32_MAX;
 }
 
+// Possibly make this faster
 uint32_t binary_search(edge_list_t *list, edge_t *elem, uint32_t start,
                        uint32_t end)
 {
@@ -260,43 +275,52 @@ class PCSR
 public:
     // data members
     std::vector<node_t> nodes;
+    std::vector<uint32_t> in_degrees;
+    std::vector<uint32_t> out_degrees;
+
     edge_list_t edges;
     uint32_t edge_count;
 
-    DEV_VEC row_offset_device;
-    DEV_VEC column_indices_device;
-    DEV_VEC eids_device;
+    // replacing device vectors
+    uint32_t *row_offset_pinned;
+    uint32_t *column_indices_pinned;
+    uint32_t *eids_pinned;
 
-    // member functions
-    PCSR(uint32_t init_n);
-    void init_graph(std::string graph_file_path);
+    // replacing device vectors
+    uint32_t *row_offset_device;
+    uint32_t *column_indices_device;
+    uint32_t *eids_device;
 
-    void add_node();
-    void add_edge(uint32_t src, uint32_t dest, uint32_t value);
-    void add_edge_update(uint32_t src, uint32_t dest, uint32_t value);
-    void edge_update_list(std::vector<std::tuple<uint32_t, uint32_t>> edge_list, bool is_delete, bool is_reverse_edge);
-    void delete_edge(uint32_t src, uint32_t dest);
+    // member functions
+    PCSR(uint32_t init_n, uint32_t max_edge_count);
+    // ~PCSR();
+    
+    // PCSR specific internal functions
     uint64_t get_n();
-    // vector<tuple<uint32_t, uint32_t, uint32_t>> get_edges();
-    void print_graph();
-    void print_array();
-    std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> get_csr_ptrs(std::vector<int> eids);
-    // void label_edges();
-    // uint32_t find_edge_id(uint32_t src, uint32_t dest);
-    // std::tuple<uint32_t, uint32_t> get_graph_attr();
-
+    void add_node();
     uint32_t insert(uint32_t index, edge_t elem, uint32_t src);
     void double_list();
     int slide_right(int index);
     void slide_left(int index);
     void fix_sentinel(int32_t node_index, int in);
     void redistribute(int index, int len);
+    void add_edge(uint32_t src, uint32_t dest, uint32_t value);
+    void add_edge_update(uint32_t src, uint32_t dest, uint32_t value);
+    void delete_edge(uint32_t src, uint32_t dest);
+
+    // exposed APIs
+    void edge_update_list(std::vector<std::tuple<uint32_t, uint32_t>> edge_list, bool is_delete, bool is_reverse_edge);
+    void label_edges();
+    void build_csr();
+    void build_reverse_csr();
+    std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> get_csr_ptrs();
+    vector<tuple<uint32_t, uint32_t, uint32_t>> get_edges();
+    void move_pinned_to_gpu();
 };
 
 ////////////////////////////////////////////////////////////////////////////
 
-// removed default argument of 0
-PCSR::PCSR(uint32_t init_n)
+PCSR::PCSR(uint32_t init_n, uint32_t max_num_edges)
 {
     if (init_n != 0)
     {
@@ -317,49 +341,34 @@ PCSR::PCSR(uint32_t init_n)
         {
             add_node();
         }
-    }
-}
-
-void PCSR::init_graph(std::string graph_file_path)
-{
-
-    FILE *fp;
-    fp = fopen(graph_file_path.c_str(), "r");
-    if (!fp)
-    {
-        printf("Open file failed.\n");
-        exit(-1);
-    }
 
-    // loading the number of nodes and edges present in
-    // the graph into the variables node_size and edge_size respectively
-    int node_size, edge_size;
-    fscanf(fp, "%d %d", &node_size, &edge_size);
+        in_degrees.resize(init_n);
+        out_degrees.resize(init_n);
 
-    edges.N = 2 << bsr_word(node_size);
-    edges.logN = (1 << bsr_word(bsr_word(edges.N) + 1));
-    edges.H = bsr_word(edges.N / edges.logN);
-    edge_count = 0;
+        cErr(cudaMallocHost(&row_offset_pinned, sizeof(uint32_t) * (init_n + 1)));
+        cErr(cudaMallocHost(&column_indices_pinned, sizeof(uint32_t) * max_num_edges));
+        cErr(cudaMallocHost(&eids_pinned, sizeof(uint32_t) * max_num_edges));
 
-    edges.items.resize(edges.N);
-    for (int i = 0; i < edges.N; i++)
-    {
-        edge_t new_edge(0, 0);
-        edges.items[i] = new_edge;
+        cErr(cudaMalloc(&row_offset_device, sizeof(uint32_t) * (init_n + 1)));
+        cErr(cudaMalloc(&column_indices_device, sizeof(uint32_t) * max_num_edges));
+        cErr(cudaMalloc(&eids_device, sizeof(uint32_t) * max_num_edges));
     }
+}
 
-    for (int i = 0; i < node_size; i++)
-    {
-        add_node();
-    }
+// PCSR::~PCSR()
+// {
+//     cErr(cudaFree(row_offset_pinned));
+//     cErr(cudaFree(column_indices_pinned));
+//     cErr(cudaFree(eids_pinned));
 
-    for (int i = 0; i < edge_size; i++)
-    {
-        int source, dest, edge_value;
-        fscanf(fp, "%d %d %d", &source, &dest, &edge_value);
-        add_edge(source, dest, edge_value);
-    }
-}
+//     cErr(cudaFree(row_offset_device));
+//     cErr(cudaFree(column_indices_device));
+//     cErr(cudaFree(eids_device));
+
+//     nodes.clear();
+//     in_degrees.clear();
+//     out_degrees.clear();
+// }
 
 // add a node to the graph
 void PCSR::add_node()
@@ -708,235 +717,152 @@ uint64_t PCSR::get_n()
     return nodes.size();
 }
 
-// vector<tuple<uint32_t, uint32_t, uint32_t>> PCSR::get_edges()
-// {
-//     uint64_t n = get_n();
-//     vector<tuple<uint32_t, uint32_t, uint32_t>> output;
-
-//     output.resize(edge_count);
-//     int iter = 0;
-//     for (int i = 0; i < n; i++)
-//     {
-//         uint32_t start = nodes[i].beginning;
-//         uint32_t end = nodes[i].end;
-//         for (int j = start + 1; j < end; j++)
-//         {
-//             if (!is_null(edges.items[j]))
-//             {
-//                 output[iter] =
-//                     make_tuple(i, edges.items[j].dest, edges.items[j].value);
-//                 iter += 1;
-//             }
-//         }
-//     }
-//     return output;
-// }
-
-void PCSR::print_graph()
+vector<tuple<uint32_t, uint32_t, uint32_t>> PCSR::get_edges()
 {
-    int num_vertices = nodes.size();
-
-    // printing the graph matrix column indices
-    for (int i = 0; i < num_vertices; ++i)
-        printf("   %d", i);
-
-    printf("\n");
+    uint64_t n = get_n();
+    vector<tuple<uint32_t, uint32_t, uint32_t>> output;
 
-    for (int i = 0; i < num_vertices; i++)
+    output.resize(edge_count);
+    int iter = 0;
+    for (int i = 0; i < n; i++)
     {
-        // +1 to avoid sentinel
-
-        // printing the graph matrix row indices
-        printf("%d ", i);
-        int matrix_index = 0;
-
-        for (uint32_t j = nodes[i].beginning + 1; j < nodes[i].end; j++)
+        uint32_t start = nodes[i].beginning;
+        uint32_t end = nodes[i].end;
+        for (int j = start + 1; j < end; j++)
         {
             if (!is_null(edges.items[j]))
             {
-                while (matrix_index < edges.items[j].dest)
-                {
-                    printf("    ");
-                    matrix_index++;
-                }
-                // printf("%03d ", edges.items[j].value);
-                printf(" %d  ", edges.items[j].value);
-                matrix_index++;
+                output[iter] =
+                    make_tuple(i, edges.items[j].dest, edges.items[j].value);
+                iter += 1;
             }
         }
-        for (uint32_t j = matrix_index; j < num_vertices; j++)
-        {
-            printf("    ");
-        }
-        printf("\n");
     }
+    return output;
 }
 
-void PCSR::print_array()
+void PCSR::label_edges()
 {
-    for (int i = 0; i < edges.N; i++)
+    uint32_t column_indices_size = edges.items.size();
+    uint32_t counter = 1;
+    for (int i = 0; i < column_indices_size; ++i)
     {
-        if (is_null(edges.items[i]))
-        {
-            printf("%d-x ", i);
-        }
-        else if (is_sentinel(edges.items[i]))
-        {
-            uint32_t value = edges.items[i].value;
-            if (value == UINT32_MAX)
-            {
-                value = 0;
-            }
-            printf("\n%d-s(%u):(%d, %d) ", i, value, nodes[value].beginning,
-                   nodes[value].end);
-        }
-        else
+        if (!is_sentinel(edges.items[i]) && !is_null(edges.items[i]))
         {
-            printf("%d-(%d, %u) ", i, edges.items[i].dest, edges.items[i].value);
+            edges.items[i].value = counter;
+            ++counter;
         }
     }
-    printf("\n\n");
 }
 
-std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> PCSR::get_csr_ptrs(std::vector<int> eids)
+void PCSR::edge_update_list(std::vector<std::tuple<uint32_t, uint32_t>> edge_list, bool is_delete = false, bool is_reverse_edge = false)
 {
+    bool is_reverse_edge_local = is_reverse_edge;
+    bool is_delete_local = is_delete;
 
-    // if (edge_count != eids.size())
-    // {
-    //     cout << "🍎🍎🍎 The eids(" << eids.size() << ") and edge_count(" << edge_count << ") [col_indices.size() = " << edges.items.size() << "] DONT match\n"
-    //          << flush;
-    // }
-
-    // we are building a compressed CSR arrays without
-    // the -1 values which indicates an empty edge
-    thrust::host_vector<int> row_offset;
-    thrust::host_vector<int> column_indices;
-
-    int row_offset_size = nodes.size() + 1;
-    int column_indices_size = edges.items.size();
-
-    // first element of row offset is always zero
-    row_offset.resize(row_offset_size);
-    column_indices.resize(edge_count);
-    row_offset[0] = 0;
-
-    for (int i = 0; i < row_offset_size - 1; ++i)
+    for (auto &edge : edge_list)
     {
-        row_offset[i + 1] = nodes[i].num_neighbors + row_offset[i];
-    }
+        uint32_t src = (is_reverse_edge_local == true) ? std::get<1>(edge) : std::get<0>(edge);
+        uint32_t dst = (is_reverse_edge_local == true) ? std::get<0>(edge) : std::get<1>(edge);
 
-    int iter = 0;
-    for (int i = 0; i < column_indices_size; ++i)
-    {
-        if (!is_sentinel(edges.items[i]) && !is_null(edges.items[i]))
-        {
-            column_indices[iter] = edges.items[i].dest;
-            iter += 1;
+        if (is_delete_local){
+            in_degrees[dst] -= 1;
+            out_degrees[src] -= 1;
+            delete_edge(src, dst);
+        }else{
+            in_degrees[dst] += 1;
+            out_degrees[src] += 1;
+            add_edge(src, dst, 1);
         }
     }
-
-    row_offset_device = row_offset;
-    column_indices_device = column_indices;
-    eids_device.resize(edge_count);
-
-    if (eids.size() == 0)
-        thrust::sequence(eids_device.begin(), eids_device.end());
-    else
-        thrust::copy(eids.begin(), eids.end(), eids_device.begin());
-
-    std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> t;
-    std::get<0>(t) = (std::uintptr_t)RAW_PTR(row_offset_device);
-    std::get<1>(t) = (std::uintptr_t)RAW_PTR(column_indices_device);
-    std::get<2>(t) = (std::uintptr_t)RAW_PTR(eids_device);
-    return t;
 }
 
-// void PCSR::label_edges()
-// {
-
-//     uint32_t column_indices_size = edges.items.size();
-//     uint32_t counter = 1;
-//     for (int i = 0; i < column_indices_size; ++i)
-//     {
-//         if (!is_sentinel(edges.items[i]) && !is_null(edges.items[i]))
-//         {
-//             edges.items[i].value = counter;
-//             ++counter;
-//         }
-//     }
-// }
-
-// uint32_t PCSR::find_edge_id(uint32_t src, uint32_t dest)
-// {
-//     edge_t e;
-//     e.dest = dest;
-//     // setting a random number
-//     e.value = 0;
-
-//     uint32_t loc = binary_search(&edges, &e, nodes[src].beginning + 1, nodes[src].end);
-//     return edges.items[loc].value;
-// }
-
-// std::tuple<uint32_t, uint32_t> PCSR::get_graph_attr()
-// {
-//     std::tuple<uint32_t, uint32_t> t;
-//     std::get<0>(t) = nodes.size();
-//     std::get<1>(t) = edge_count;
-//     return t;
-// }
-
-// void copy_label_edges(PCSR &pcsr, PCSR &ref_pcsr)
-// {
-//     uint64_t n = pcsr.get_n();
-
-//     for (int i = 0; i < n; i++)
-//     {
-//         uint32_t start = pcsr.nodes[i].beginning;
-//         uint32_t end = pcsr.nodes[i].end;
-//         for (int j = start + 1; j < end; j++)
-//         {
-//             if (!is_null(pcsr.edges.items[j]))
-//             {
-//                 // Searching for the edge_id of reverse edge from reference_pcsr
-//                 pcsr.edges.items[j].value = ref_pcsr.find_edge_id(pcsr.edges.items[j].dest, i);
-//             }
-//         }
-//     }
-// }
-
-void PCSR::edge_update_list(std::vector<std::tuple<uint32_t, uint32_t>> edge_list, bool is_delete = false, bool is_reverse_edge = false)
+void PCSR::build_reverse_csr()
 {
-    // cout << "📦📦📦 Edge Update list" << endl
-    //      << flush;
-    for (auto &edge : edge_list)
+    uint64_t n = get_n();
+    // computing the bwd row offsets
+    row_offset_pinned[0] = in_degrees[0];
+    for(int i=1; i<in_degrees.size(); ++i){
+        row_offset_pinned[i] = row_offset_pinned[i-1] + in_degrees[i];
+    }
+    row_offset_pinned[in_degrees.size()] = edge_count;
+    
+    for (int i = 0; i < n; i++)
     {
-        uint32_t src = (is_reverse_edge == true) ? std::get<1>(edge) : std::get<0>(edge);
-        uint32_t dst = (is_reverse_edge == true) ? std::get<0>(edge) : std::get<1>(edge);
+        uint32_t start = nodes[i].beginning;
+        uint32_t end = nodes[i].end;
+        for (int j = start + 1; j < end; j++)
+        {
+            if (!is_sentinel(edges.items[j]) && !is_null(edges.items[j]))
+            {
+                row_offset_pinned[edges.items[j].dest] -= 1;
 
-        if (is_delete)
-            delete_edge(src, dst);
-        else
-            add_edge(src, dst, 1);
+                int col_index = row_offset_pinned[edges.items[j].dest];
+                column_indices_pinned[col_index] = i;
+                eids_pinned[col_index] = edges.items[j].value;
+            }
+        }
     }
+
+    move_pinned_to_gpu();
 }
 
-void build_reverse_pcsr(PCSR &pcsr, PCSR &ref_pcsr)
+void PCSR::build_csr()
 {
-    // cout << "💄💄💄 Building Reverse PCSR" << endl
-    //      << flush;
-    uint64_t n = ref_pcsr.get_n();
+    // computing the bwd row offsets
+    uint64_t n = get_n();
+    row_offset_pinned[0] = out_degrees[0];
+    for(int i=1; i<out_degrees.size(); ++i){
+        row_offset_pinned[i] = row_offset_pinned[i-1] + out_degrees[i];
+    }
+    row_offset_pinned[out_degrees.size()] = edge_count;
+    
     for (int i = 0; i < n; i++)
     {
-        uint32_t start = ref_pcsr.nodes[i].beginning;
-        uint32_t end = ref_pcsr.nodes[i].end;
+        uint32_t start = nodes[i].beginning;
+        uint32_t end = nodes[i].end;
         for (int j = start + 1; j < end; j++)
         {
-            if (!is_sentinel(ref_pcsr.edges.items[j]) && !is_null(ref_pcsr.edges.items[j]))
+            if (!is_sentinel(edges.items[j]) && !is_null(edges.items[j]))
             {
-                pcsr.add_edge(ref_pcsr.edges.items[j].dest, i, ref_pcsr.edges.items[j].value);
+                row_offset_pinned[i] -= 1;
+                column_indices_pinned[row_offset_pinned[i]] = edges.items[j].dest;
+                eids_pinned[row_offset_pinned[i]] = edges.items[j].value;
             }
         }
     }
+
+    move_pinned_to_gpu();
+}
+
+void PCSR::move_pinned_to_gpu(){
+    cErr(cudaMemcpy(row_offset_device, row_offset_pinned, sizeof(uint32_t) * (get_n() + 1), cudaMemcpyHostToDevice));
+    cErr(cudaMemcpy(column_indices_device, column_indices_pinned, sizeof(uint32_t) * edge_count, cudaMemcpyHostToDevice));
+    cErr(cudaMemcpy(eids_device, eids_pinned, sizeof(uint32_t) * edge_count, cudaMemcpyHostToDevice));
+}
+
+std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> PCSR::get_csr_ptrs(){
+    std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> t;
+    std::get<0>(t) = (std::uintptr_t)row_offset_device;
+    std::get<1>(t) = (std::uintptr_t)column_indices_device;
+    std::get<2>(t) = (std::uintptr_t)eids_device;
+    return t;
+}
+
+std::vector<std::vector<uint32_t>> read_gpu_csr(PCSR &pcsr){
+    std::vector<uint32_t> row_offset(pcsr.get_n() + 1);
+    std::vector<uint32_t> column_indices(pcsr.edge_count);
+    std::vector<uint32_t> eids(pcsr.edge_count);
+
+    cErr(cudaMemcpy(row_offset.data(), pcsr.row_offset_device, sizeof(uint32_t) * (pcsr.get_n() + 1), cudaMemcpyDeviceToHost));
+    cErr(cudaMemcpy(column_indices.data(), pcsr.column_indices_device, sizeof(uint32_t) * pcsr.edge_count, cudaMemcpyDeviceToHost));
+    cErr(cudaMemcpy(eids.data(), pcsr.eids_device, sizeof(uint32_t) * pcsr.edge_count, cudaMemcpyDeviceToHost));
+
+    std::vector<std::vector<uint32_t>> res;
+    res.push_back(row_offset);
+    res.push_back(column_indices);
+    res.push_back(eids);
+    return res;
 }
 
 ////////////////////////////////////////////////////////////////////////////
@@ -946,63 +872,20 @@ void build_reverse_pcsr(PCSR &pcsr, PCSR &ref_pcsr)
 PYBIND11_MODULE(pcsr, m)
 {
     m.doc() = "PCSR Dynamic Graph Representation";
-    // m.def("copy_label_edges", &copy_label_edges, "Label edges of a PCSR based on another PCSR");
-    m.def("build_reverse_pcsr", &build_reverse_pcsr, "Builds the reverse PCSR based on another PCSR");
-
-    py::class_<node_t>(m, "Node")
-        .def(py::init<int, int, int>(), py::arg("beg") = 0, py::arg("_end") = 0, py::arg("num_neigh") = 0)
-        .def_readwrite("beginning", &_node::beginning)
-        .def_readwrite("end", &_node::end)
-        .def_readwrite("num_neighbors", &_node::num_neighbors)
-        .def_readwrite("in_degree", &_node::in_degree);
-
-    py::class_<edge_t>(m, "Edge")
-        .def(py::init<int, int>(), py::arg("_dest") = 0, py::arg("_value") = 0)
-        .def_readwrite("dest", &edge_t::dest)
-        .def_readwrite("value", &edge_t::value);
-
-    py::class_<edge_list_t>(m, "EdgeList")
-        .def(py::init<>())
-        .def_readwrite("N", &edge_list_t::N)
-        .def_readwrite("H", &edge_list_t::H)
-        .def_readwrite("logN", &edge_list_t::logN)
-        .def_readwrite("items", &edge_list_t::items);
-
-    py::class_<pair_int>(m, "PairInt")
-        .def(py::init<int, int>())
-        .def_readwrite("x", &pair_int::x)
-        .def_readwrite("y", &pair_int::y);
-
-    py::class_<pair_double>(m, "PairDouble")
-        .def(py::init<double, double>())
-        .def_readwrite("x", &pair_double::x)
-        .def_readwrite("y", &pair_double::y);
+    m.def("read_gpu_csr", &read_gpu_csr, "Read CSR arrays from the GPU");
 
     py::class_<PCSR>(m, "PCSR")
-        .def(py::init<int>(), py::arg("init_n") = 0)
-        .def("init_graph", &PCSR::init_graph)
-        .def("add_node", &PCSR::add_node)
-        .def("insert", &PCSR::insert)
-        .def("double_list", &PCSR::double_list)
-        .def("slide_right", &PCSR::slide_right)
-        .def("slide_left", &PCSR::slide_left)
-        .def("fix_sentinel", &PCSR::fix_sentinel)
-        .def("redistribute", &PCSR::redistribute)
-        .def("print_graph", &PCSR::print_graph)
-        .def("add_edge", &PCSR::add_edge)
-        .def("add_edge_update", &PCSR::add_edge_update)
-        .def("edge_update_list", &PCSR::edge_update_list, py::arg("edge_list"), py::arg("is_delete") = false, py::arg("is_reverse_edge") = false)
-        // .def("label_edges", &PCSR::label_edges, "Creates edge labels for the current GPMA")
-        .def("delete_edge", &PCSR::delete_edge)
-        .def("get_n", &PCSR::get_n)
-        // .def("get_edges", &PCSR::get_edges)
-        .def("print_array", &PCSR::print_array)
-        .def("get_csr_ptrs", &PCSR::get_csr_ptrs, py::arg("eids"))
-        // .def("get_graph_attr", &PCSR::get_graph_attr)
-        // .def("find_edge_id", &PCSR::find_edge_id)
-        .def_readwrite("nodes", &PCSR::nodes)
-        .def_readwrite("edges", &PCSR::edges)
+        .def(py::init<int, int>(), py::arg("init_n"), py::arg("max_edge_count"))
+        .def_readwrite("in_degrees", &PCSR::in_degrees)
+        .def_readwrite("out_degrees", &PCSR::out_degrees)
         .def_readwrite("edge_count", &PCSR::edge_count)
+        .def("get_n", &PCSR::get_n)
+        .def("edge_update_list", &PCSR::edge_update_list, py::arg("edge_list"), py::arg("is_delete") = false, py::arg("is_reverse_edge") = false)
+        .def("label_edges", &PCSR::label_edges, "Creates edge labels for the current GPMA")
+        .def("get_edges", &PCSR::get_edges)
+        .def("build_csr", &PCSR::build_csr)
+        .def("build_reverse_csr", &PCSR::build_reverse_csr)
+        .def("get_csr_ptrs", &PCSR::get_csr_ptrs)
         .def("__copy__", [](const PCSR &self)
              { return PCSR(self); })
         .def(
@@ -1011,4 +894,4 @@ PYBIND11_MODULE(pcsr, m)
             "memo"_a);
 }
 
-// /usr/local/cuda-11.7/bin/nvcc $(python3 -m pybind11 --includes) -shared -rdc=true --compiler-options '-fPIC'  -D__CDPRT_SUPPRESS_SYNC_DEPRECATION_WARNING -o pcsr.so pcsr.cu
\ No newline at end of file
+// /usr/local/cuda-11.7/bin/nvcc $(python3 -m pybind11 --includes) -shared -rdc=true --compiler-options '-fPIC'  -D__CDPRT_SUPPRESS_SYNC_DEPRECATION_WARNING -o pcsr.so pcsr.cu
diff --git a/seastar/graph/static/StaticGraph.py b/seastar/graph/static/StaticGraph.py
index 0a5485f..b0cea4a 100644
--- a/seastar/graph/static/StaticGraph.py
+++ b/seastar/graph/static/StaticGraph.py
@@ -13,29 +13,22 @@
 from seastar.graph.static.csr import CSR
 
 class StaticGraph(SeastarGraph):
-    def __init__(self, edge_list):    
+    def __init__(self, edge_list, edge_weights, num_nodes):    
         super().__init__()
-        self._num_nodes = 0
-        self._num_edges = 0
-        
-        console.log("Getting graph attributes")
-        
-        self._get_graph_attr(edge_list)
-        
-        console.log("Building forward edge list")
+        self._num_nodes = num_nodes
+        self._num_edges = len(set(edge_list))
         
+        # console.log("Building forward edge list")
         self._prepare_edge_lst_fwd(edge_list)
+        # console.log("Creating forward graph")
+        self._forward_graph = CSR(self.fwd_edge_list, edge_weights, self._num_nodes, is_edge_reverse=True)
         
-        console.log("Building backward edge list")
-        
-        self._prepare_edge_lst_bwd(self.fwd_edge_list)  
-        
-        console.log("Creating forward graph")
-        self._forward_graph = CSR(self.fwd_edge_list, self._num_nodes, is_edge_reverse=True)
-        
-        console.log("Creating backward graph")
-        self._backward_graph = CSR(self.bwd_edge_list, self._num_nodes)
+        # console.log("Building backward edge list")
+        self._prepare_edge_lst_bwd(self.fwd_edge_list)
+        # console.log("Creating backward graph")
+        self._backward_graph = CSR(self.bwd_edge_list, edge_weights, self._num_nodes)
         
+        # console.log("Getting CSR ptrs")
         self._get_graph_csr_ptrs()
     
     def _prepare_edge_lst_fwd(self, edge_list):    
@@ -45,21 +38,10 @@ def _prepare_edge_lst_fwd(self, edge_list):
         self.fwd_edge_list = edge_list_for_t
     
     def _prepare_edge_lst_bwd(self, edge_list):    
-        edge_list_for_t = edge_list
+        edge_list_for_t = copy.deepcopy(edge_list)
         edge_list_for_t.sort()
         self.bwd_edge_list = edge_list_for_t
         
-    def _get_graph_attr(self, edge_list):
-        node_set = set()
-        for edge in edge_list:
-            src, dst = edge[0], edge[1]
-            node_set.add(src)
-            node_set.add(dst)
-        
-        self._num_nodes = len(node_set)
-        self._num_edges = len(edge_list)
-            
-        
     def _get_graph_csr_ptrs(self):
         fwd_csr_ptrs = self._forward_graph
         self.fwd_row_offset_ptr = fwd_csr_ptrs.row_offset_ptr
@@ -76,6 +58,15 @@ def get_num_nodes(self):
     
     def get_num_edges(self):
         return self._num_edges
+    
+    def get_ndata(self, field):
+        if field in self._ndata:
+            return self._ndata[field]
+        else:
+            return None
+
+    def set_ndata(self, field, val):
+        self._ndata[field] = val
         
     def graph_type(self):
         return "csr"
@@ -84,4 +75,7 @@ def in_degrees(self):
         return np.array(self._forward_graph.out_degrees, dtype='int32')
     
     def out_degrees(self):
-        return np.array(self._forward_graph.in_degrees, dtype='int32')
\ No newline at end of file
+        return np.array(self._forward_graph.in_degrees, dtype='int32')
+    
+    def weighted_in_degrees(self):
+        return np.array(self._forward_graph.weighted_out_degrees, dtype='int32')
\ No newline at end of file
diff --git a/seastar/graph/static/csr.cu b/seastar/graph/static/csr.cu
index 6cf5bfb..00e5068 100644
--- a/seastar/graph/static/csr.cu
+++ b/seastar/graph/static/csr.cu
@@ -32,12 +32,13 @@ public:
 
     std::vector<int> in_degrees;
     std::vector<int> out_degrees;
+    std::vector<float> weighted_out_degrees;
 
     std::uintptr_t row_offset_ptr;
     std::uintptr_t column_indices_ptr;
     std::uintptr_t eids_ptr;
 
-    CSR(std::vector<std::tuple<int, int, int>> edge_list, int num_nodes, bool is_edge_reverse);
+    CSR(std::vector<std::tuple<int, int, int>> edge_list, std::vector<float> edge_weight, int num_nodes, bool is_edge_reverse);
     // std::tuple<std::uintptr_t, std::uintptr_t, std::uintptr_t> get_csr_ptrs();
     void get_csr_ptrs();
     // void label_edges();
@@ -55,7 +56,7 @@ bool sort_by_sec(const std::tuple<int, int> &a,
            std::tie(std::get<1>(b), std::get<0>(b));
 }
 
-CSR::CSR(std::vector<std::tuple<int, int, int>> edge_list, int num_nodes, bool is_edge_reverse = false)
+CSR::CSR(std::vector<std::tuple<int, int, int>> edge_list, std::vector<float> edge_weight, int num_nodes, bool is_edge_reverse = false)
 {
 
     // initialising row_offset values all to -1
@@ -67,6 +68,7 @@ CSR::CSR(std::vector<std::tuple<int, int, int>> edge_list, int num_nodes, bool i
 
     in_degrees.resize(num_nodes, 0);
     out_degrees.resize(num_nodes, 0);
+    weighted_out_degrees.resize(num_nodes, 0);
 
     int current_src;
     int beg = 0;
@@ -103,6 +105,7 @@ CSR::CSR(std::vector<std::tuple<int, int, int>> edge_list, int num_nodes, bool i
         // updating the degree arrays
         out_degrees[src] += 1;
         in_degrees[dst] += 1;
+        weighted_out_degrees[src] += edge_weight[eid];
     }
 
     row_offset[current_src + 1] = end;
@@ -235,14 +238,24 @@ void print_dev_array(std::uintptr_t ptr, int size)
     std::cout << "\n";
 }
 
+std::vector<int> get_array(std::uintptr_t ptr, int size)
+{
+    int *dev_ptr = reinterpret_cast<int *>(ptr);
+    int *host_ptr = (int *)malloc(size * sizeof(int));
+    // thrust::device_ptr<int> dev_ptr = thrust::device_pointer_cast(raw_ptr);
+    cudaMemcpy(host_ptr, dev_ptr, size * sizeof(int), cudaMemcpyDeviceToHost);
+    std::vector<int> vec(host_ptr,host_ptr+size);
+    return vec;
+}
 PYBIND11_MODULE(csr, m)
 {
     m.doc() = "CPython module for CSR"; // optional module docstring
     m.def("print_dev_array", &print_dev_array, "Print arrays on the device");
+    m.def("get_array",&get_array,"Get Array");
 
     py::class_<CSR>(m, "CSR")
         // .def(py::init<std::vector<std::tuple<int, int>>, int, bool>(), py::arg("edge_list"), py::arg("num_nodes"), py::arg("is_edge_reverse") = false)
-        .def(py::init<std::vector<std::tuple<int, int, int>>, int, bool>(), py::arg("edge_list"), py::arg("num_nodes"), py::arg("is_edge_reverse") = false)
+        .def(py::init<std::vector<std::tuple<int, int, int>>, std::vector<float>, int, bool>(), py::arg("edge_list"), py::arg("edge_weight"), py::arg("num_nodes"), py::arg("is_edge_reverse") = false)
         // .def("label_edges", &CSR::label_edges)
         // .def("copy_label_edges", &CSR::copy_label_edges)
         .def("get_csr_ptrs", &CSR::get_csr_ptrs)
@@ -255,6 +268,7 @@ PYBIND11_MODULE(csr, m)
         .def_readwrite("row_offset_ptr", &CSR::row_offset_ptr)
         .def_readwrite("column_indices_ptr", &CSR::column_indices_ptr)
         .def_readwrite("eids_ptr", &CSR::eids_ptr)
+        .def_readwrite("weighted_out_degrees", &CSR::weighted_out_degrees)
         .def_readwrite("out_degrees", &CSR::out_degrees)
         .def_readwrite("in_degrees", &CSR::in_degrees)
         .def("__copy__", [](const CSR &self)
diff --git a/seastar/nn/pytorch/__init__.py b/seastar/nn/pytorch/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/seastar/nn/pytorch/gat_conv.py b/seastar/nn/pytorch/gat_conv.py
new file mode 100644
index 0000000..af71e34
--- /dev/null
+++ b/seastar/nn/pytorch/gat_conv.py
@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+from seastar.compiler.backend.pytorch.torch_callback import SeastarBackendTorch
+from seastar.compiler import Seastar
+
+# pylint: enable=W0235
+class GATConv(nn.Module):
+    def __init__(self,
+                 in_feats,
+                 out_feats,
+                 num_heads,
+                 feat_drop=0.,
+                 attn_drop=0.,
+                 negative_slope=0.2,
+                 activation=None):
+        super(GATConv, self).__init__()
+        self._num_heads = num_heads
+        self._in_feats = in_feats
+        self._out_feats = out_feats
+        self.fc = nn.Linear(
+            self._in_feats, out_feats * num_heads, bias=False)
+        self.attn_l = nn.Parameter(torch.FloatTensor(size=(num_heads, out_feats)))
+        self.attn_r = nn.Parameter(torch.FloatTensor(size=(num_heads, out_feats)))
+        self.feat_drop = nn.Dropout(feat_drop)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.leaky_relu = nn.LeakyReLU(negative_slope)
+        self.negative_slope = negative_slope
+        
+        self.activation = activation
+        self.seastar = Seastar(SeastarBackendTorch())
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """Reinitialize learnable parameters."""
+        gain = nn.init.calculate_gain('relu')
+        nn.init.xavier_normal_(self.fc.weight, gain=gain)
+        nn.init.xavier_normal_(self.attn_l, gain=gain)
+        nn.init.xavier_normal_(self.attn_r, gain=gain)
+
+    def forward(self, graph, feat):
+
+        h_dst = h_src = self.feat_drop(feat)
+        feat_src = feat_dst = self.fc(h_src).view(-1, self._num_heads, self._out_feats)
+        el = (feat_src * self.attn_l).sum(dim=-1).unsqueeze(-1)
+        er = (feat_dst * self.attn_r).sum(dim=-1).unsqueeze(-1)
+
+        # Vertex-centric implementation.
+        @self.seastar.compile(gnn_module=self)
+        def nb_forward(v):
+            embs = [nb.el + v.er for nb in v.innbs]
+            coeff = [torch.exp(self.leaky_relu(emb - max(embs))) for emb in embs]
+            s = sum(coeff)
+            alpha = [c/s for c in coeff]
+            feat_src = [nb.feat_src for nb in v.innbs]
+            return sum([alpha[i] * feat_src[i] for i in range(len(feat_src))])
+        rst = nb_forward(g=graph, n_feats= {'el':el, 'er': er, 'feat_src':feat_src})
+
+        # activation
+        if self.activation:
+            rst = self.activation(rst)
+        return rst
diff --git a/seastar/nn/pytorch/graph_conv.py b/seastar/nn/pytorch/graph_conv.py
new file mode 100644
index 0000000..a7be3dd
--- /dev/null
+++ b/seastar/nn/pytorch/graph_conv.py
@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+from seastar.compiler import Seastar
+from seastar.compiler.backend.pytorch.torch_callback import SeastarBackendTorch
+
+class GraphConv(nn.Module):
+    def __init__(self,
+                 in_feats,
+                 out_feats,
+                 activation,
+                 bias=True):
+        super(GraphConv, self).__init__()
+        self.weight = nn.Parameter(torch.Tensor(in_feats, out_feats))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_feats))
+        else:
+            self.bias = None
+        self.activation = activation
+        self.seastar = Seastar(SeastarBackendTorch())
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.weight)
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def forward(self, g, h, edge_weight=None):
+        h = torch.mm(h, self.weight)
+        
+        if edge_weight is None:
+            @self.seastar.compile(gnn_module=self)
+            def nb_compute(v):
+                h = sum([nb.h*nb.norm for nb in v.innbs])
+                h = h * v.norm
+                return h
+            h = nb_compute(g=g, n_feats={'norm': g.get_ndata("norm"), 'h' : h})
+        else:
+            @self.seastar.compile(gnn_module=self)
+            def nb_compute(v):
+                h = sum([nb_edge.src.norm * nb_edge.src.h * nb_edge.edge_weight for nb_edge in v.inedges])
+                h = h * v.norm
+                return h
+            h = nb_compute(g=g, n_feats={'norm': g.get_ndata("norm"), 'h' : h}, e_feats={'edge_weight':edge_weight})
+
+        # bias
+        if self.bias is not None:
+            h = h + self.bias
+        if self.activation:
+            h = self.activation(h)
+        return h
diff --git a/seastar/nn/pytorch/temporal/__init__.py b/seastar/nn/pytorch/temporal/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/seastar/nn/pytorch/temporal/tgcn.py b/seastar/nn/pytorch/temporal/tgcn.py
new file mode 100644
index 0000000..2fe1ab9
--- /dev/null
+++ b/seastar/nn/pytorch/temporal/tgcn.py
@@ -0,0 +1,72 @@
+import torch
+from seastar.nn.pytorch.graph_conv import GraphConv
+
+
+class TGCN(torch.nn.Module):
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+    ):
+        super(TGCN, self).__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        # Update GCN Layer
+        self.conv_z = GraphConv(self.in_channels, self.out_channels, activation=None)
+       
+        # Update linear layer
+        self.linear_z = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Reset GCN layer
+        self.conv_r = GraphConv(self.in_channels, self.out_channels, activation=None)
+
+        # Reset linear layer
+        self.linear_r = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+        # Candidate (Current Memory Content) GCN layer
+        self.conv_h = GraphConv(self.in_channels, self.out_channels, activation=None)
+
+        # Candidate linear layer
+        self.linear_h = torch.nn.Linear(2 * self.out_channels, self.out_channels)
+
+    def _set_hidden_state(self, X, H):
+        if H is None:
+            H = torch.zeros(X.shape[0], self.out_channels).to(X.device)
+        return H
+
+    def _calculate_update_gate(self, g, X, edge_weight, H):
+        h = self.conv_z(g, X, edge_weight=edge_weight)
+        Z = torch.cat((h, H), dim=1)
+        Z = self.linear_z(Z)
+        Z = torch.sigmoid(Z)
+        return Z
+
+    def _calculate_reset_gate(self, g, X, edge_weight, H):
+        h = self.conv_r(g, X, edge_weight=edge_weight)
+        R = torch.cat((h, H), dim=1)
+        R = self.linear_r(R)
+        R = torch.sigmoid(R)
+        return R
+
+    def _calculate_candidate_state(self, g, X, edge_weight, H, R):
+        h = self.conv_h(g, X, edge_weight=edge_weight)
+        H_tilde = torch.cat((h, H * R), dim=1)
+        H_tilde = self.linear_h(H_tilde)
+        H_tilde = torch.tanh(H_tilde)
+        return H_tilde
+
+    def _calculate_hidden_state(self, Z, H, H_tilde):
+        H = Z * H + (1 - Z) * H_tilde
+        return H
+
+    def forward(self, g, X, edge_weight=None, H=None):
+
+        H = self._set_hidden_state(X, H)
+        Z = self._calculate_update_gate(g, X, edge_weight, H)
+        R = self._calculate_reset_gate(g, X, edge_weight, H)
+        H_tilde = self._calculate_candidate_state(g, X, edge_weight, H, R)
+        H = self._calculate_hidden_state(Z, H, H_tilde)
+        return H


Test patch:
diff --git a/profiler/test_add.cu b/profiler/test_add.cu
deleted file mode 100644
index 7f6a03e..0000000
--- a/profiler/test_add.cu
+++ /dev/null
@@ -1,41 +0,0 @@
-#include <thrust/device_vector.h>
-#include <thrust/host_vector.h>
-#include <thrust/remove.h>
-#include <thrust/sort.h>
-#include <cub/cub.cuh>
-
-int main()
-{
-    thrust::host_vector<int> vec_1(500, 10);
-
-    cudaEvent_t start, stop;
-    cudaEventCreate(&start);
-    cudaEventCreate(&stop);
-
-    cudaEventRecord(start);
-    thrust::device_vector<int> vec_2 = vec_1;
-    cudaEventRecord(stop);
-
-    cudaEventSynchronize(stop);
-    float milliseconds = 0;
-    cudaEventElapsedTime(&milliseconds, start, stop);
-
-    printf("\n⏱️ Time taken: %.3f\n", milliseconds);
-
-    thrust::host_vector<int> vec_3(500, 10);
-    thrust::device_vector<int> d_vec_3(500);
-    cudaEventCreate(&start);
-    cudaEventCreate(&stop);
-
-    cudaEventRecord(start);
-    for (int i = 0; i < 500; ++i)
-    {
-        d_vec_3[i] = vec_3[i];
-    }
-    cudaEventRecord(stop);
-
-    cudaEventSynchronize(stop);
-    cudaEventElapsedTime(&milliseconds, start, stop);
-
-    printf("\n⏱️ Time taken: %.3f\n", milliseconds);
-}
\ No newline at end of file
diff --git a/profiler/test_profiler.py b/profiler/test_profiler.py
deleted file mode 100644
index 1f65b84..0000000
--- a/profiler/test_profiler.py
+++ /dev/null
@@ -1,8 +0,0 @@
-import cProfile
-
-def factorial(n):
-    if n == 1:
-        return 1
-    return n * factorial(n-1)
-
-cProfile.run('factorial(500)')
\ No newline at end of file
diff --git a/profiler/test_vec_add.cu b/profiler/test_vec_add.cu
deleted file mode 100644
index e068f23..0000000
--- a/profiler/test_vec_add.cu
+++ /dev/null
@@ -1,131 +0,0 @@
-/*
- * A simplified example of vector addition in CUDA to illustrate the
- * data decomposition pattern using blocks of threads.
- *
- * To compile:
- *   nvcc -o va-GPU-simple VA-GPU-simple.cu
- */
-
-#include <stdio.h>
-
-// In this example we use a very small number of blocks
-// and threads in those blocks for illustration
-// on a very small array
-#define N 100
-#define numThread 2 // 2 threads in a block
-#define numBlock 4  // 4 blocks
-
-/*
- * 1.
- *  The 'kernel' function that will be executed on the GPU device hardware.
- */
-__global__ void add(int *a, int *b, int *c)
-{
-
-    // the initial index that this thread will work on
-    int tid = blockDim.x * blockIdx.x + threadIdx.x;
-
-    // In this above example code, we assume a linear set of blocks of threads in the 'x' dimension,
-    // which is declared in main below when we run this function.
-
-    // The actual computation is being done by individual threads
-    // in each of the blocks.
-    // e.g. we use 4 blocks and 2 threads per block (8 threads will run in parallel)
-    //      and our total array size N is 8
-    //      the thread whose threadIdx.x is 0 within block 0 will compute c[0],
-    //          because tid = (2 * 0)  + 0
-    //      the thread whose threadIdx.x is 0 within block 1 will compute c[2],
-    //          because tid = (2 * 1) + 0
-    //      the thread whose threadIdx.x is 1 within block 1 will compute c[3],
-    //          because tid = (2 * 1) + 1
-    //
-    //     The following while loop will execute once for this simple example:
-    //          c[0] through c[7] will be computed concurrently
-    //
-    while (tid < N)
-    {
-        c[tid] = a[tid] + b[tid]; // The actual computation done by the thread
-        tid += blockDim.x;        // Increment this thread's index by the number of threads per block:
-                                  // in this small case, each thread would then have a tid > N
-    }
-}
-
-/*
- * The main program that directs the execution of vector add on the GPU
- */
-int main(void)
-{
-    int *a, *b, *c;             // The arrays on the host CPU machine
-    int *dev_a, *dev_b, *dev_c; // The arrays for the GPU device
-
-    // 2.a allocate the memory on the CPU
-    a = (int *)malloc(N * sizeof(int));
-    b = (int *)malloc(N * sizeof(int));
-    c = (int *)malloc(N * sizeof(int));
-
-    // 2.b. fill the arrays 'a' and 'b' on the CPU with dummy values
-    for (int i = 0; i < N; i++)
-    {
-        a[i] = i;
-        b[i] = i;
-    }
-
-    // 2.c. allocate the memory on the GPU
-    cudaMalloc((void **)&dev_a, N * sizeof(int));
-    cudaMalloc((void **)&dev_b, N * sizeof(int));
-    cudaMalloc((void **)&dev_c, N * sizeof(int));
-
-    // 2.d. copy the arrays 'a' and 'b' to the GPU
-    cudaMemcpy(dev_a, a, N * sizeof(int),
-               cudaMemcpyHostToDevice);
-    cudaMemcpy(dev_b, b, N * sizeof(int),
-               cudaMemcpyHostToDevice);
-
-    cudaEvent_t start, stop;
-    cudaEventCreate(&start);
-    cudaEventCreate(&stop);
-
-    // 3. Execute the vector addition 'kernel function' on th GPU device,
-    // declaring how many blocks and how many threads per block to use.
-    cudaEventRecord(start);
-    add<<<numBlock, numThread>>>(dev_a, dev_b, dev_c);
-    cudaEventRecord(stop);
-
-    // 4. copy the array 'c' back from the GPU to the CPU
-    cudaMemcpy(c, dev_c, N * sizeof(int),
-               cudaMemcpyDeviceToHost);
-
-    cudaEventSynchronize(stop);
-    float milliseconds = 0;
-    cudaEventElapsedTime(&milliseconds, start, stop);
-
-    printf("\n⏱️ Time taken: %f\n", milliseconds);
-
-    // verify that the GPU did the work we requested
-    bool success = true;
-    int total = 0;
-    printf("Checking %d values in the array.\n", N);
-    for (int i = 0; i < N; i++)
-    {
-        if ((a[i] + b[i]) != c[i])
-        {
-            printf("Error:  %d + %d != %d\n", a[i], b[i], c[i]);
-            success = false;
-        }
-        total += 1;
-    }
-    if (success)
-        printf("We did it, %d values correct!\n", total);
-
-    // free the memory we allocated on the CPU
-    free(a);
-    free(b);
-    free(c);
-
-    // free the memory we allocated on the GPU
-    cudaFree(dev_a);
-    cudaFree(dev_b);
-    cudaFree(dev_c);
-
-    return 0;
-}
\ No newline at end of file
diff --git a/tests/graph/dynamic/pcsr/test_PCSRGraph.py b/tests/graph/dynamic/pcsr/test_PCSRGraph.py
deleted file mode 100644
index 7be4d37..0000000
--- a/tests/graph/dynamic/pcsr/test_PCSRGraph.py
+++ /dev/null
@@ -1,5 +0,0 @@
-"""
-    Unit tests for seastar.graph.dynamic.pcsr.PCSRGraph module
-"""
-
-from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
\ No newline at end of file
diff --git a/tests/graph/dynamic/pcsr/test_pcsr.py b/tests/graph/dynamic/pcsr/test_pcsr.py
deleted file mode 100644
index 58c0409..0000000
--- a/tests/graph/dynamic/pcsr/test_pcsr.py
+++ /dev/null
@@ -1,10 +0,0 @@
-"""
-    Unit tests for seastar.graph.dynamic.pcsr.pcsr.PCSR class
-"""
-
-from seastar.graph.dynamic.pcsr.pcsr import PCSR
-
-class TestPCSR:
-    
-    def test_init(self):
-        pcsr = PCSR()
\ No newline at end of file
diff --git a/tests/graph/dynamic/test_DynamicGraph.py b/tests/graph/dynamic/test_DynamicGraph.py
deleted file mode 100644
index fa69bbf..0000000
--- a/tests/graph/dynamic/test_DynamicGraph.py
+++ /dev/null
@@ -1,102 +0,0 @@
-import json
-from rich import inspect
-from seastar.graph.dynamic.DynamicGraph import DynamicGraph
-from seastar.graph.dynamic.naive.NaiveGraph import NaiveGraph
-
-
-def get_edge_list(gra):
-    pass
-
-
-class TestDynamicGraph:
-    # Stores the edge-list for each timestamp
-    # edge_index[t] has the edge list for timestamp = t
-    edge_index = [
-        [[0, 1], [1, 2], [2, 3]],
-        [[1, 2], [2, 3], [3, 1]],
-        [[1, 2], [2, 3], [3, 1], [3, 0]],
-        [[2, 3], [3, 1], [3, 0], [0, 1]],
-    ]
-
-    # stores the edge-weights for each timestamp
-    # edge_weight[t] has the edge weight for timestamp = t
-    # corresponding to the edge in edge_index
-    edge_weight = [[4, 7, 9], [2, 11, 13], [3, 8, 1, 5], [15, 6, 10, 12]]
-
-    # total timestamps for this dynamic graph dataset
-    time_periods = 4
-
-    # node features for each timestamp
-    y = [[2, 11, 15, 8], [9, 5, 7, 10], [1, 3, 17, 19], [4, 6, 12, 13]]
-
-    def test_get_num_nodes(self):
-        """
-            Assert the number of nodes in the graph, then repeat
-            this assert step after calling DynamicGraph.get_graph()
-            and DynamicGraph.get_backward_graph() in sequential order
-        """
-
-        # base graph: t = 0
-        naive_graph = NaiveGraph(edge_list=self.edge_index)
-        assert naive_graph.get_num_nodes() == 4
-
-        # graph: t = 1
-        naive_graph.get_graph(1)
-        assert naive_graph.get_num_nodes() == 3
-
-        # graph: t = 2
-        naive_graph.get_graph(2)
-        assert naive_graph.get_num_nodes() == 4
-
-        # graph: t = 3
-        naive_graph.get_graph(3)
-        assert naive_graph.get_num_nodes() == 4
-
-        # Now moving the graph in the backward direction
-        # graph: t = 2
-        naive_graph.get_backward_graph(2)
-        assert naive_graph.get_num_nodes() == 4
-        
-        # graph: t = 1
-        naive_graph.get_backward_graph(1)
-        assert naive_graph.get_num_nodes() == 3
-        
-        # graph: t = 1
-        naive_graph.get_backward_graph(0)
-        assert naive_graph.get_num_nodes() == 4
-        
-    def test_get_num_edges(self):
-        """
-            Assert the number of edges in the graph, then repeat
-            this assert step after calling DynamicGraph.get_graph()
-            and DynamicGraph.get_backward_graph() in sequential order
-        """
-        
-        # base graph: t = 0
-        naive_graph = NaiveGraph(edge_list=self.edge_index)
-        assert naive_graph.get_num_edges() == 3
-
-        # graph: t = 1
-        naive_graph.get_graph(1)
-        assert naive_graph.get_num_edges() == 3
-
-        # graph: t = 2
-        naive_graph.get_graph(2)
-        assert naive_graph.get_num_edges() == 4
-
-        # graph: t = 3
-        naive_graph.get_graph(3)
-        assert naive_graph.get_num_edges() == 4
-
-        # Now moving the graph in the backward direction
-        # graph: t = 2
-        naive_graph.get_backward_graph(2)
-        assert naive_graph.get_num_edges() == 4
-        
-        # graph: t = 1
-        naive_graph.get_backward_graph(1)
-        assert naive_graph.get_num_edges() == 3
-        
-        # graph: t = 1
-        naive_graph.get_backward_graph(0)
-        assert naive_graph.get_num_edges() == 3
\ No newline at end of file
diff --git a/tests/graph/dynamic/test_DynamicGraphPCSR.py b/tests/graph/dynamic/test_DynamicGraphPCSR.py
new file mode 100644
index 0000000..24f4ea4
--- /dev/null
+++ b/tests/graph/dynamic/test_DynamicGraphPCSR.py
@@ -0,0 +1,278 @@
+import json
+from rich import inspect
+
+from seastar.graph.dynamic.pcsr.PCSRGraph import PCSRGraph
+from seastar.graph.static.csr import get_dev_array
+
+class TestDynamicGraphPCSR:
+    # Stores the edge-list for each timestamp
+    # edge_index[t] has the edge list for timestamp = t
+    edge_index = [
+        [[0, 1], [1, 2], [2, 3]],
+        [[1, 2], [2, 3], [3, 1]],
+        [[1, 2], [2, 3], [3, 1], [3, 0]],
+        [[2, 3], [3, 1], [3, 0], [0, 1]],
+    ]
+
+    # sorted based on second-element, then first-element
+    sorted_edge_index = [
+        [[0, 1], [1, 2], [2, 3]],
+        [[3, 1], [1, 2], [2, 3]],
+        [[3, 0], [3, 1], [1, 2], [2, 3]],
+        [[3, 0], [0, 1], [3, 1], [2, 3]],
+    ]
+
+    # stores the edge-weights for each timestamp
+    # edge_weight[t] has the edge weight for timestamp = t
+    # corresponding to the edge in edge_index
+    edge_weight = [[4, 7, 9], [2, 11, 13], [3, 8, 1, 5], [15, 6, 10, 12]]
+
+    # total timestamps for this dynamic graph dataset
+    time_periods = 4
+
+    # node features for each timestamp
+    y = [[2, 11, 15, 8], [9, 5, 7, 10], [1, 3, 17, 19], [4, 6, 12, 13]]
+
+    def test_get_graph_attr(self):
+        pcsr_graph = PCSRGraph(edge_list=self.sorted_edge_index, max_num_nodes=4)
+        graph_attr = pcsr_graph._get_graph_attr(edge_list=self.sorted_edge_index)
+
+        # checking if the size of graph_attr = 4, which is the
+        # total number of timestamps present
+        assert len(graph_attr) == 4
+
+        # checking the (num_nodes, num_edges) pair for each timestamp
+        assert graph_attr["0"] == (4, 3)
+        assert graph_attr["1"] == (4, 3)
+        assert graph_attr["2"] == (4, 4)
+        assert graph_attr["3"] == (4, 4)
+
+    def test_preprocess_graph_structure(self):
+        pcsr_graph = PCSRGraph(edge_list=self.sorted_edge_index, max_num_nodes=4)
+
+        # checking graph_updates for t = 0
+        graph_updates_0 = pcsr_graph.graph_updates["0"]
+        assert len(graph_updates_0["add"]) == 3
+        for edge in [(0, 1), (1, 2), (2, 3)]:
+            assert edge in graph_updates_0["add"]
+
+        assert len(graph_updates_0["delete"]) == 0
+
+        assert graph_updates_0["num_nodes"] == 4
+        assert graph_updates_0["num_edges"] == 3
+
+        # checking graph_updates for t = 1
+        graph_updates_1 = pcsr_graph.graph_updates["1"]
+        assert len(graph_updates_1["add"]) == 1
+        assert (3, 1) in graph_updates_1["add"]
+
+        assert len(graph_updates_1["delete"]) == 1
+        assert (0, 1) in graph_updates_1["delete"]
+
+        assert graph_updates_1["num_nodes"] == 4
+        assert graph_updates_1["num_edges"] == 3
+
+        # checking graph_updates for t = 2
+        graph_updates_2 = pcsr_graph.graph_updates["2"]
+
+        assert len(graph_updates_2["add"]) == 1
+        assert (3, 0) in graph_updates_2["add"]
+
+        assert len(graph_updates_2["delete"]) == 0
+
+        assert graph_updates_2["num_nodes"] == 4
+        assert graph_updates_2["num_edges"] == 4
+
+        # checking graph_updates for t = 3
+        graph_updates_3 = pcsr_graph.graph_updates["3"]
+
+        assert len(graph_updates_3["add"]) == 1
+        assert (0, 1) in graph_updates_3["add"]
+
+        assert len(graph_updates_3["delete"]) == 1
+        assert (1, 2) in graph_updates_3["delete"]
+
+        assert graph_updates_3["num_nodes"] == 4
+        assert graph_updates_3["num_edges"] == 4
+
+    def test_get_graph(self):
+        pcsr_graph = PCSRGraph(edge_list=self.sorted_edge_index, max_num_nodes=4)
+
+        # for time = 0
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.fwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.fwd_column_indices_ptr, 3),
+            get_dev_array(pcsr_graph.fwd_eids_ptr, 3),
+        )
+
+        assert row_offset == [0, 0, 1, 2, 3]
+        assert column_indices == [0, 1, 2]
+        assert eids == [0, 1, 2]
+
+        # for time = 1
+        pcsr_graph.get_graph(1)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.fwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.fwd_column_indices_ptr, 3),
+            get_dev_array(pcsr_graph.fwd_eids_ptr, 3),
+        )
+
+        assert row_offset == [0, 0, 1, 2, 3]
+        assert column_indices == [3, 1, 2]
+        assert eids == [0, 1, 2]
+
+        # for time = 2
+        pcsr_graph.get_graph(2)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.fwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.fwd_column_indices_ptr, 4),
+            get_dev_array(pcsr_graph.fwd_eids_ptr, 4),
+        )
+
+        assert row_offset == [0, 1, 2, 3, 4]
+        assert column_indices == [3, 3, 1, 2]
+        assert eids == [0, 1, 2, 3]
+
+        # for time = 3
+        pcsr_graph.get_graph(3)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.fwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.fwd_column_indices_ptr, 4),
+            get_dev_array(pcsr_graph.fwd_eids_ptr, 4),
+        )
+
+        assert row_offset == [0, 1, 3, 3, 4]
+        assert column_indices == [3, 0, 3, 2]
+        assert eids == [0, 1, 2, 3]
+
+    def test_get_backward_graph(self):
+        
+        pcsr_graph = PCSRGraph(edge_list=self.sorted_edge_index, max_num_nodes=4)
+        pcsr_graph.get_graph(3)
+
+        # for time = 3
+        pcsr_graph.get_backward_graph(3)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.bwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.bwd_column_indices_ptr, 4),
+            get_dev_array(pcsr_graph.bwd_eids_ptr, 4),
+        )
+        
+        assert row_offset == [0,1,1,2,4]
+        assert column_indices == [1,3,0,1]
+        assert eids == [1,3,0,2]
+        
+        # for time = 2
+        pcsr_graph.get_backward_graph(2)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.bwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.bwd_column_indices_ptr, 4),
+            get_dev_array(pcsr_graph.bwd_eids_ptr, 4),
+        )
+        
+        assert row_offset == [0,0,1,2,4]
+        assert column_indices == [2,3,0,1]
+        assert eids == [2,3,0,1]
+        
+        # for time = 1
+        pcsr_graph.get_backward_graph(1)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.bwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.bwd_column_indices_ptr, 3),
+            get_dev_array(pcsr_graph.bwd_eids_ptr, 3),
+        )
+        
+        assert row_offset == [0,0,1,2,3]
+        assert column_indices == [2,3,1]
+        assert eids == [1,2,0]
+        
+        # for time = 0
+        pcsr_graph.get_backward_graph(0)
+        
+        row_offset, column_indices, eids = (
+            get_dev_array(pcsr_graph.bwd_row_offset_ptr, 5),
+            get_dev_array(pcsr_graph.bwd_column_indices_ptr, 3),
+            get_dev_array(pcsr_graph.bwd_eids_ptr, 3),
+        )
+        
+        assert row_offset == [0,1,2,3,3]
+        assert column_indices == [1,2,3]
+        assert eids == [0,1,2]
+
+    def test_get_num_nodes(self):
+        """
+        Assert the number of nodes in the graph, then repeat
+        this assert step after calling DynamicGraph.get_graph()
+        and DynamicGraph.get_backward_graph() in sequential order
+        """
+
+        # base graph: t = 0
+        pcsr_graph = PCSRGraph(edge_list=self.edge_index, max_num_nodes=4)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # graph: t = 1
+        pcsr_graph.get_graph(1)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # graph: t = 2
+        pcsr_graph.get_graph(2)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # graph: t = 3
+        pcsr_graph.get_graph(3)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # Now moving the graph in the backward direction
+        # graph: t = 2
+        pcsr_graph.get_backward_graph(2)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # graph: t = 1
+        pcsr_graph.get_backward_graph(1)
+        assert pcsr_graph.get_num_nodes() == 4
+
+        # graph: t = 1
+        pcsr_graph.get_backward_graph(0)
+        assert pcsr_graph.get_num_nodes() == 4
+
+    def test_get_num_edges(self):
+        """
+        Assert the number of edges in the graph, then repeat
+        this assert step after calling DynamicGraph.get_graph()
+        and DynamicGraph.get_backward_graph() in sequential order
+        """
+
+        # base graph: t = 0
+        pcsr_graph = PCSRGraph(edge_list=self.edge_index, max_num_nodes=4)
+        assert pcsr_graph.get_num_edges() == 3
+
+        # graph: t = 1
+        pcsr_graph.get_graph(1)
+        assert pcsr_graph.get_num_edges() == 3
+
+        # graph: t = 2
+        pcsr_graph.get_graph(2)
+        assert pcsr_graph.get_num_edges() == 4
+
+        # graph: t = 3
+        pcsr_graph.get_graph(3)
+        assert pcsr_graph.get_num_edges() == 4
+
+        # Now moving the graph in the backward direction
+        # graph: t = 2
+        pcsr_graph.get_backward_graph(2)
+        assert pcsr_graph.get_num_edges() == 4
+
+        # graph: t = 1
+        pcsr_graph.get_backward_graph(1)
+        assert pcsr_graph.get_num_edges() == 3
+
+        # graph: t = 1
+        pcsr_graph.get_backward_graph(0)
+        assert pcsr_graph.get_num_edges() == 3
\ No newline at end of file