diff --git a/pkg/go/graph/weighted_graph.go b/pkg/go/graph/weighted_graph.go
index 65e395dc..8ebaa9af 100644
--- a/pkg/go/graph/weighted_graph.go
+++ b/pkg/go/graph/weighted_graph.go
@@ -14,18 +14,32 @@ type WeightedAuthorizationModelGraph struct {
 	drawingDirection DrawingDirection
 }
 
+//nolint: cyclop
 func NewWeightedAuthorizationModelGraph(model *openfgav1.AuthorizationModel) (*WeightedAuthorizationModelGraph, error) {
 	g, err := NewAuthorizationModelGraph(model)
 	if err != nil {
 		return nil, err
 	}
-	g, err = g.Reversed() // we want direction of Check
+	g, err = g.Reversed() // we want edges to have the direction of Check when doing the weight assignments later
 	if err != nil {
 		return nil, err
 	}
 
 	graphBuilder := &WeightedAuthorizationModelGraphBuilder{multi.NewDirectedGraph()}
 
+	// Add all nodes
+	iterNodes := g.Nodes()
+	for iterNodes.Next() {
+		nextNode := iterNodes.Node()
+		from, ok := nextNode.(*AuthorizationModelNode)
+		if !ok {
+			return nil, fmt.Errorf("%w: could not cast to WeightedAuthorizationModelNode", ErrBuildingGraph)
+		}
+		newNode := &WeightedAuthorizationModelNode{from, make(WeightMap), false}
+		graphBuilder.AddNode(newNode)
+	}
+
+	// Add all the edges
 	iterEdges := g.Edges()
 	for iterEdges.Next() {
 		nextEdge, ok := iterEdges.Edge().(multi.Edge)
diff --git a/pkg/go/graph/weighted_graph_builder.go b/pkg/go/graph/weighted_graph_builder.go
index fcacc134..d5a30f55 100644
--- a/pkg/go/graph/weighted_graph_builder.go
+++ b/pkg/go/graph/weighted_graph_builder.go
@@ -13,29 +13,9 @@ type WeightedAuthorizationModelGraphBuilder struct {
 }
 
 func (wb *WeightedAuthorizationModelGraphBuilder) AddEdgeWithWeights(edge *AuthorizationModelEdge) error {
-	isLoop := edge.From().ID() == edge.To().ID()
-
-	// create "from" node
 	fromNode := wb.Node(edge.From().ID())
-	if fromNode == nil {
-		from, ok := edge.From().(*AuthorizationModelNode)
-		if !ok {
-			return fmt.Errorf("%w: could not cast to WeightedAuthorizationModelNode", ErrBuildingGraph)
-		}
-		fromNode = &WeightedAuthorizationModelNode{from, make(WeightMap), false}
-		wb.AddNode(fromNode)
-	}
-
-	// create "to" node
 	toNode := wb.Node(edge.To().ID())
-	if toNode == nil {
-		to, ok := edge.To().(*AuthorizationModelNode)
-		if !ok {
-			return fmt.Errorf("%w: could not cast to WeightedAuthorizationModelNode", ErrBuildingGraph)
-		}
-		toNode = &WeightedAuthorizationModelNode{to, make(WeightMap), false}
-		wb.AddNode(toNode)
-	}
+	isLoop := fromNode == toNode
 
 	// update "isNested" field for the node
 	if isLoop {
diff --git a/pkg/go/graph/weighted_graph_builder_test.go b/pkg/go/graph/weighted_graph_builder_test.go
index 818e6505..e9e260d8 100644
--- a/pkg/go/graph/weighted_graph_builder_test.go
+++ b/pkg/go/graph/weighted_graph_builder_test.go
@@ -31,6 +31,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#viewer - weights:[user=1]"];
 2 [label=user];
 
@@ -53,6 +54,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#viewer - weights:[user1=1,user2=1]"];
 2 [label=user1];
 3 [label=user2];
@@ -79,6 +81,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#viewer - weights:[user=1]"];
 2 [label=user];
 3 [label="user:*"];
@@ -103,6 +106,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#rewrite - weights:[user=1]"];
 2 [label="folder#viewer - weights:[user=1]"];
 3 [label=user];
@@ -130,6 +134,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#viewer - weights:[user=+∞]"];
 2 [label=user];
 
@@ -159,6 +164,7 @@ rankdir=TB
 0 [label=company];
 1 [label="company#approved_member - weights:[user=1]"];
 2 [label=user];
+3 [label=license];
 4 [label="license#member - weights:[user=2]"];
 5 [label="license#owner - weights:[company=1]"];
 
@@ -197,6 +203,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#org - weights:[org=1]"];
 2 [label=org];
 3 [label="folder#viewer - weights:[user=4]"];
@@ -255,6 +262,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=group];
 1 [label="group#max_owner - weights:[user=+∞]"];
 2 [label="group#owner - weights:[user=3]"];
 3 [label=user];
@@ -263,6 +271,7 @@ rankdir=TB
 6 [label="state#can_view - weights:[user=1]"];
 7 [label="union - weights:[user=1]"];
 8 [label="state#member - weights:[user=1]"];
+9 [label=transition];
 10 [label="intersection - weights:[user=2]"];
 11 [label="transition#end - weights:[state=1]"];
 12 [label="transition#start - weights:[state=1]"];
@@ -333,6 +342,7 @@ rankdir=TB
 ];
 
 // Node definitions.
+0 [label=folder];
 1 [label="folder#a - weights:[user1=1,user2=1]"];
 2 [label=user1];
 3 [label=user2];