Merge pull request #375 from WeiPhil/register_loss_fix

Added interface to register loss from outside library
NVlabs · Oct 13, 2023 · 2121041 · 2121041
2 parents bd1f727 + c305956
commit 2121041
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/include/tiny-cuda-nn/loss.h b/include/tiny-cuda-nn/loss.h
@@ -70,4 +70,7 @@ std::unique_ptr<Loss<T>> default_loss(const std::string& name) {
 
 std::vector<std::string> builtin_losses();
 
+template <typename T>
+void register_loss(const std::string& name, const std::function<Loss<T>*(const json&)>& factory);
+
 }
diff --git a/src/loss.cu b/src/loss.cu
@@ -78,6 +78,9 @@ void register_loss(const std::string& name, const std::function<Loss<T>*(const j
 	register_loss(loss_factories<T>(), name, factory);
 }
 
+template void register_loss<float>(const std::string& name, const std::function<Loss<float>*(const json&)>& factory);
+template void register_loss<__half>(const std::string& name, const std::function<Loss<__half>*(const json&)>& factory);
+
 template <typename T>
 Loss<T>* create_loss(const json& loss) {
 	std::string name = loss.value("otype", "RelativeL2");