add dcocumentation and export

mlverse · Aug 19, 2020 · 702091a · 702091a
1 parent 83c2bf6
commit 702091a
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 0 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -99,6 +99,7 @@ export(is_torch_dtype)
 export(is_torch_layout)
 export(is_torch_memory_format)
 export(is_torch_qscheme)
+export(load_state_dict)
 export(nn_adaptive_avg_pool1d)
 export(nn_adaptive_avg_pool2d)
 export(nn_adaptive_avg_pool3d)

diff --git a/R/save.R b/R/save.R
@@ -66,6 +66,21 @@ torch_load_module <- function(obj) {
   obj$module
 }
 
+#' Load a state dict file
+#' 
+#' This function should only be used to load models saved in python.
+#' For it to work correctly you need to use `torch.save` with the flag:
+#' `_use_new_zipfile_serialization=True` and also remove all `nn.Parameter`
+#' classes from the tensors in the dict. 
+#' 
+#' The above might change with development of [this](https://github.com/pytorch/pytorch/issues/37213) 
+#' in pytorch's C++ api.
+#' 
+#' @param path to the state dict file
+#' 
+#' @return a named list of tensors.
+#'
+#' @export
 load_state_dict <- function(path) {
   path <- normalizePath(path)
   o <- cpp_load_state_dict(path)

diff --git a/man/load_state_dict.Rd b/man/load_state_dict.Rd