Merge branch 'INSTX-3280-fix-modbase-perf' into 'master'

[INSTX-3280] Bring back macOS modbase performance See merge request machine-learning/dorado!819
nanoporetech · Jan 24, 2024 · 74b4b53 · 74b4b53
2 parents e48bfea + d75776e
commit 74b4b53
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 10 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -44,7 +44,7 @@ function(enable_warnings_as_errors TARGET_NAME)
             _CRT_SECURE_NO_WARNINGS
         )
     else()
-        target_compile_options(${TARGET_NAME} PRIVATE -Wall -Werror -Wextra)
+        target_compile_options(${TARGET_NAME} PRIVATE -Wall -Werror -Wextra -Wundef)
     endif()
 endfunction()
 

diff --git a/dorado/modbase/ModBaseCaller.cpp b/dorado/modbase/ModBaseCaller.cpp
@@ -72,17 +72,10 @@ ModBaseCaller::ModBaseCaller(const std::vector<std::filesystem::path>& model_pat
     if (device == "cpu") {
         // no slow_conv2d_cpu for type Half, need to use float32
         m_options = at::TensorOptions().device(torch::kCPU).dtype(torch::kFloat32);
+#ifdef __APPLE__
     } else if (device == "metal") {
-#if TORCH_VERSION_MAJOR < 2
-        // no metal implementation yet, force to cpu
-        auto torchMetalBackend = torch::kCPU;
-        auto torchMetalDtype = torch::kFloat32;
-        spdlog::debug("- no metal backend available for modified basecalling, defaulting to CPU.");
-#else
-        auto torchMetalBackend = torch::kMPS;
-        auto torchMetalDtype = torch::kFloat16;
+        m_options = at::TensorOptions().device(torch::kMPS).dtype(torch::kFloat16);
 #endif
-        m_options = at::TensorOptions().device(torchMetalBackend).dtype(torchMetalDtype);
     } else {
         m_options = at::TensorOptions().device(device).dtype(torch::kFloat16);
     }