diff --git a/.gitignore b/.gitignore
index c948f680..5c0c9bdf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,3 +4,4 @@ inst/doc
 .Rhistory
 .venv
 activate
+.V8history
diff --git a/DESCRIPTION b/DESCRIPTION
index 9ecf7c04..21ee5aaa 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -17,7 +17,7 @@ RoxygenNote: 7.1.1
 URL: https://github.com/mlverse/tabnet
 BugReports: https://github.com/mlverse/tabnet/issues
 Imports: 
-    torch (>= 0.3.0.9001),
+    torch (>= 0.4.0),
     hardhat,
     magrittr,
     glue,
@@ -43,5 +43,3 @@ Suggests:
     tidyr,
     purrr
 VignetteBuilder: knitr
-Remotes:
-    mlverse/torch
diff --git a/man/autoplot.tabnet_explain.Rd b/man/autoplot.tabnet_explain.Rd
index 90a7fafc..b5d63c98 100644
--- a/man/autoplot.tabnet_explain.Rd
+++ b/man/autoplot.tabnet_explain.Rd
@@ -4,7 +4,12 @@
 \alias{autoplot.tabnet_explain}
 \title{Plot tabnet_explain mask importance heatmap}
 \usage{
-\method{autoplot}{tabnet_explain}(object, type = c("mask_agg", "steps"), quantile = 1, ...)
+autoplot.tabnet_explain(
+  object,
+  type = c("mask_agg", "steps"),
+  quantile = 1,
+  ...
+)
 }
 \arguments{
 \item{object}{A \code{tabnet_explain} object as a result of \code{\link[=tabnet_explain]{tabnet_explain()}}.}
diff --git a/man/autoplot.tabnet_fit.Rd b/man/autoplot.tabnet_fit.Rd
index 934e8f67..da978636 100644
--- a/man/autoplot.tabnet_fit.Rd
+++ b/man/autoplot.tabnet_fit.Rd
@@ -5,9 +5,9 @@
 \alias{autoplot.tabnet_pretrain}
 \title{Plot tabnet_fit model loss along epochs}
 \usage{
-\method{autoplot}{tabnet_fit}(object, ...)
+autoplot.tabnet_fit(object, ...)
 
-\method{autoplot}{tabnet_pretrain}(object, ...)
+autoplot.tabnet_pretrain(object, ...)
 }
 \arguments{
 \item{object}{A \code{tabnet_fit} or \code{tabnet_pretrain} object as a result of