diff --git a/src/MLOps.NET/MLOps.NET.csproj b/src/MLOps.NET/MLOps.NET.csproj
index 288fd820..f6025b9f 100644
--- a/src/MLOps.NET/MLOps.NET.csproj
+++ b/src/MLOps.NET/MLOps.NET.csproj
@@ -23,7 +23,7 @@
     <PackageReference Include="Dynamitey" Version="2.0.10.189" />
     <PackageReference Include="ICSharpCode.Decompiler" Version="7.1.0.6543" />
     <PackageReference Include="Microsoft.Extensions.DependencyModel" Version="5.0.0" />
-    <PackageReference Include="Microsoft.ML" Version="1.6.0" />
+    <PackageReference Include="Microsoft.ML" Version="2.0.1" />
     <PackageReference Include="System.IO.Abstractions" Version="13.2.47" />
     <PackageReference Include="Microsoft.EntityFrameworkCore" Version="5.0.11" />
     <PackageReference Include="Microsoft.EntityFrameworkCore.Tools" Version="5.0.11"><IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
diff --git a/test/MLOps.NET.IntegrationTests/MLOps.NET.IntegrationTests.csproj b/test/MLOps.NET.IntegrationTests/MLOps.NET.IntegrationTests.csproj
index cda7c701..333bf606 100644
--- a/test/MLOps.NET.IntegrationTests/MLOps.NET.IntegrationTests.csproj
+++ b/test/MLOps.NET.IntegrationTests/MLOps.NET.IntegrationTests.csproj
@@ -8,7 +8,7 @@
     <PackageReference Include="FluentAssertions" Version="6.2.0" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.0.0" />
     <PackageReference Include="Moq" Version="4.16.1" />
-    <PackageReference Include="Microsoft.ML" Version="1.6.0" />
+    <PackageReference Include="Microsoft.ML" Version="2.0.1" />
     <PackageReference Include="MSTest.TestAdapter" Version="2.2.7" />
     <PackageReference Include="MSTest.TestFramework" Version="2.2.7" />
     <PackageReference Include="coverlet.collector" Version="3.1.0"><IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
diff --git a/test/MLOps.NET.Tests.Common/MLOps.NET.Tests.Common.csproj b/test/MLOps.NET.Tests.Common/MLOps.NET.Tests.Common.csproj
index 58c18a8a..b5623dbb 100644
--- a/test/MLOps.NET.Tests.Common/MLOps.NET.Tests.Common.csproj
+++ b/test/MLOps.NET.Tests.Common/MLOps.NET.Tests.Common.csproj
@@ -10,7 +10,7 @@
     <PackageReference Include="Microsoft.Extensions.Configuration.Abstractions" Version="6.0.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration.EnvironmentVariables" Version="5.0.0" />
     <PackageReference Include="Microsoft.Extensions.Configuration.Json" Version="5.0.0" />
-    <PackageReference Include="Microsoft.ML" Version="1.6.0" />
+    <PackageReference Include="Microsoft.ML" Version="2.0.1" />
   </ItemGroup>
 
   <ItemGroup>
diff --git a/test/MLOps.NET.Tests/MLOps.NET.Tests.csproj b/test/MLOps.NET.Tests/MLOps.NET.Tests.csproj
index 6169fa62..172098b9 100644
--- a/test/MLOps.NET.Tests/MLOps.NET.Tests.csproj
+++ b/test/MLOps.NET.Tests/MLOps.NET.Tests.csproj
@@ -7,7 +7,7 @@
   <ItemGroup>
     <PackageReference Include="FluentAssertions" Version="6.2.0" />
     <PackageReference Include="Microsoft.EntityFrameworkCore.InMemory" Version="5.0.11" />
-    <PackageReference Include="Microsoft.ML" Version="1.6.0" />
+    <PackageReference Include="Microsoft.ML" Version="2.0.1" />
     <PackageReference Include="Microsoft.ML.FastTree" Version="1.6.0" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" Version="17.0.0" />
     <PackageReference Include="MSTest.TestAdapter" Version="2.2.7" />