Merge pull request #22 from normal-computing/remove-item

Remove .item in favour of .detach

.gitignore

@@ -1,6 +1,7 @@
 _archive/
 .vscode/
 *.pkl
+*.DS_Store
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

uqlib/ekf/diag_fisher.py

@@ -113,7 +113,7 @@ def update(
         grad,
         inplace=inplace,
     )
-    return EKFDiagState(update_mean, update_sd_diag, log_liks.mean().item(), aux)
+    return EKFDiagState(update_mean, update_sd_diag, log_liks.mean().detach(), aux)
 
 
 def build(

uqlib/ekf/diag_hessian.py

@@ -86,7 +86,7 @@ def update(
         grad,
         inplace=inplace,
     )
-    return EKFDiagState(update_mean, update_sd_diag, log_lik.item(), aux)
+    return EKFDiagState(update_mean, update_sd_diag, log_lik.detach(), aux)
 
 
 def build(

uqlib/sgmcmc/sghmc.py

@@ -20,7 +20,7 @@ class SGHMCState(NamedTuple):
 
     params: TensorTree
     momenta: TensorTree
-    log_posterior: float = 0.0
+    log_posterior: torch.tensor = torch.tensor(0.0)
     aux: Any = None
 
 
@@ -90,7 +90,7 @@ def transform_momenta(m, g):
     )
     momenta = flexi_tree_map(transform_momenta, state.momenta, grads, inplace=inplace)
 
-    return SGHMCState(params, momenta, log_post.item(), aux)
+    return SGHMCState(params, momenta, log_post.detach(), aux)
 
 
 def build(

uqlib/vi/diag.py

@@ -118,7 +118,7 @@ def nelbo_log_sd(m, lsd):
     mean, log_sd_diag = torchopt.apply_updates(
         (state.mean, state.log_sd_diag), updates, inplace=inplace
     )
-    return VIDiagState(mean, log_sd_diag, optimizer_state, nelbo_val.item(), aux)
+    return VIDiagState(mean, log_sd_diag, optimizer_state, nelbo_val.detach(), aux)
 
 
 def build(