diff --git a/lib/Scheduler/Scheduler.cpp b/lib/Scheduler/Scheduler.cpp
index 25923fffdf..dd1a37892e 100644
--- a/lib/Scheduler/Scheduler.cpp
+++ b/lib/Scheduler/Scheduler.cpp
@@ -22,19 +22,19 @@ bool isTTIROp(mlir::Operation *op) {
 }
 
 bool isTTShedulableOp(mlir::Operation *op) {
-  return isTTNNOp(op) || isTTIROp(op);
+  return (isTTNNOp(op) || isTTIROp(op)) && (not isa<func::ReturnOp>(op));
 }
 
 // Init the dependencies map of all ops which are TTIR ops
 Scheduler::Scheduler(func::FuncOp *func) {
-  for (auto &op : func->getOps()) {
+  for (mlir::Operation &op : func->getOps()) {
     if (isTTShedulableOp(&op)) {
       dependencies[&op] = {};
       unscheduledOps.insert(&op);
     }
   }
 
-  for (auto &op : func->getOps()) {
+  for (mlir::Operation &op : func->getOps()) {
     // Skip non TTIR operations
     // Skip operations which do not implement DestinationStyleOpInterface
     if (!isTTShedulableOp(&op)) {
@@ -46,7 +46,7 @@ Scheduler::Scheduler(func::FuncOp *func) {
     for (mlir::Operation *use : result.getUsers()) {
       // Skip non TTIR operations
       // Skip operations which set the result
-      if (isTTShedulableOp(use) && use->getResult(0) != result) {
+      if (isTTShedulableOp(use)) {
         dependencies[use].push_back(&op);
       }
     }