diff --git a/pyproject.toml b/pyproject.toml
index 2e748695..73dca1bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -66,7 +66,6 @@ only-include = ["pyproject.toml", "README.md"]
 [tool.basedpyright] 
 typeCheckingMode = "strict"
 failOnWarnings = true 
-stubPath = "stubs"
 
 reportAny = "error"
 reportUnknownVariableType = "error"
diff --git a/shared/protobufs/types/mlx/nn/__init__.pyi b/shared/protobufs/types/mlx/nn/__init__.pyi
new file mode 100644
index 00000000..464c4f1a
--- /dev/null
+++ b/shared/protobufs/types/mlx/nn/__init__.pyi
@@ -0,0 +1,3 @@
+from mlx.nn.layers import *
+from mlx.nn import init as init, losses as losses
+from mlx.nn.utils import average_gradients as average_gradients, value_and_grad as value_and_grad
\ No newline at end of file
diff --git a/shared/types/worker/commands_runner.py b/shared/types/worker/commands_runner.py
index 5a4b8cfa..57d66fd7 100644
--- a/shared/types/worker/commands_runner.py
+++ b/shared/types/worker/commands_runner.py
@@ -11,7 +11,7 @@ from shared.openai import FinishReason
 
 class MessageType(str, Enum):
     Setup = 'setup'
-    ChatTask = "chat_task"
+    ChatTask = 'chat_task'
     Exit = 'exit'
 
 MT = TypeVar(name='MT', bound=MessageType)
diff --git a/shared/types/worker/resource_monitor.py b/shared/types/worker/resource_monitor.py
index 7ac27b2f..ccb115f3 100644
--- a/shared/types/worker/resource_monitor.py
+++ b/shared/types/worker/resource_monitor.py
@@ -3,11 +3,11 @@ from collections.abc import Coroutine
 
 import asyncio
 
-from shared.types.events.events import ResourceProfiledEvent
+from shared.types.events.events import ResourceProfiled
 from shared.types.profiling.common import NodePerformanceProfile, MemoryPerformanceProfile, SystemPerformanceProfile
 
 class EventLog:
-    def append(self, event: ResourceProfiledEvent) -> None:
+    def append(self, event: ResourceProfiled) -> None:
         ...
 
 class ResourceCollector(ABC):
diff --git a/shared/utils.py b/shared/utils.py
index 3a9acce1..da09cb04 100644
--- a/shared/utils.py
+++ b/shared/utils.py
@@ -2,7 +2,7 @@ from typing import Any, Type, TypeVar
 
 T = TypeVar('T')
 
-def ensure_type(obj: Any, expected_type: Type[T]) -> T:
+def ensure_type(obj: Any, expected_type: Type[T]) -> T: # type: ignore
     if not isinstance(obj, expected_type):
-        raise TypeError(f"Expected {expected_type}, got {type(obj)}")
+        raise TypeError(f"Expected {expected_type}, got {type(obj)}") # type: ignore
     return obj
\ No newline at end of file