lstm cuda test

Author: liugangcode

docs/source/api/utils.rst

@@ -26,14 +26,34 @@ Utility Functions
    :undoc-members:
    :show-inheritance:
 
-.. automodule:: torch_molecule.utils.graph
+.. automodule:: torch_molecule.utils.graph.features
    :members:
    :undoc-members:
    :show-inheritance:
 
-.. automodule:: torch_molecule.utils.generic
+.. automodule:: torch_molecule.utils.graph.graph_from_smiles
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: torch_molecule.utils.graph.graph_to_smiles
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: torch_molecule.utils.generic.metrics
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+.. automodule:: torch_molecule.utils.generic.pseudo_tasks
    :members:
    :undoc-members:
    :show-inheritance:
 
-   
+.. automodule:: torch_molecule.utils.generic.weights
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
+

docs/source/conf.py

@@ -24,7 +24,6 @@
 templates_path = ['_templates']
 exclude_patterns = []
 
-
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 

tests/predictor/run_lstm.py

@@ -1,7 +1,7 @@
 import numpy as np
 from torch_molecule import LSTMMolecularPredictor
 from torch_molecule.utils.search import ParameterType, ParameterSpec
-
+import torch
 def test_lstm_predictor():
     # Test data
     smiles_list = [
@@ -67,13 +67,15 @@ def test_lstm_predictor():
     model.save_to_local(save_path)
     print(f"Model saved to {save_path}")
 
+    print(f"Device: {torch.device('cuda' if torch.cuda.is_available() else 'cpu')}")
     new_model = LSTMMolecularPredictor(
         task_type="regression",
         output_dim=15,
         LSTMunits=60,
         batch_size=2,
         epochs=2,
-        device="cpu"
+        # device="cpu"
+        device=torch.device("cuda" if torch.cuda.is_available() else "cpu")
     )
     new_model.load_from_local(save_path)
     print("Model loaded successfully")