Repository: mindsdb/lightwood

Issue:
Improve documentation of device-related methods and docs regarding GPU usage
At the moment the policy to trigger GPU usage (and advanced features like AMP) is not clear when reading the documentation.


Discussion:
I want more information about what exactly I will write in this section? @paxcema 
I'm not sure this is an ideal "good first issue". Regardless, I'd recommend that you take a close look at both the `lightwood.helpers.device` module and the `lightwood.helpers.torch.LightwoodAutocast` module, then search for their usage throughout the repo.

The ideal Pull Request would add informative docstrings to the classes and methods above where missing, explaining how, when and why are GPUs leveraged when using Lightwood.
I would like to work on this issue...I can do it

Hey @paxcema , Can you give a detailed explanation what exactly is to be done? 
Hi, I think the comment above summarizes it. Essentially, take a look at the code and document what it is doing as you understand it (doesn't have to be perfect, we can adjust it in the PR process).

Resolving patch:
diff --git a/lightwood/helpers/device.py b/lightwood/helpers/device.py
--- a/lightwood/helpers/device.py
+++ b/lightwood/helpers/device.py
@@ -5,6 +5,24 @@
 
 
 def is_cuda_compatible():
+    """
+    Check if the system has CUDA-compatible devices with the required architecture and
+    compiled CUDA version.
+
+    This function checks the compatibility of CUDA devices available on the system by
+    comparing their architectures and the compiled CUDA version. It iterates through
+    the available devices and verifies if their architectures meet the minimum
+    requirement specified by the function, and also checks if the compiled CUDA
+    version is greater than a specific version.
+
+    Returns:
+        bool: True if there are compatible CUDA devices, otherwise False.
+
+    Example:
+        >>> is_compatible = is_cuda_compatible()
+        >>> print(is_compatible)
+        True
+    """
     compatible_device_count = 0
     if torch.version.cuda is not None:
         for d in range(device_count()):
@@ -23,6 +41,27 @@ def is_cuda_compatible():
 
 
 def get_devices():
+    """
+    Get the appropriate Torch device(s) based on CUDA availability and compatibility.
+
+    This function determines the appropriate Torch device(s) to be used for
+    computations based on the availability of CUDA and compatible devices. It checks
+    if CUDA is available and if the available CUDA devices are compatible according to
+    the 'is_cuda_compatible()' function. If compatible devices are found, the function
+    selects either the first available CUDA device or a randomly selected one based on
+    the 'RANDOM_GPU' environment variable. If CUDA is not available or no compatible
+    devices are found, the function returns the CPU device.
+
+    Returns:
+        Tuple: A tuple containing the selected Torch device and the number of available
+        devices.
+    Example:
+        >>> device, num_devices = get_devices()
+        >>> print(device)
+        cuda:0
+        >>> print(num_devices)
+        1
+    """
     if torch.cuda.is_available() and is_cuda_compatible():
         device_str = "cuda"
         available_devices = torch.cuda.device_count()
@@ -40,10 +79,22 @@ def get_devices():
 
 def get_device_from_name(device_name=''):
     """
-    Returns the device specified as an argument.
-    If the argument is left empty it will returns the output of get_devices().
-    
-    :param device_name: name of the device to use (default is an empty string), if is an empty string will use the output of get_devices() instead")
+    Get a Torch device based on the specified device name or default behavior.
+
+    This function returns a Torch device based on the specified device name or the
+    default behavior, which is to return the output of the 'get_devices()' function.
+
+    Args:
+        device_name (str, optional): Name of the device to use. Default is an empty
+        string.
+        
+    Returns:
+        torch.device: The selected Torch device.
+
+    Example:
+        >>> device = get_device_from_name('cuda:1')
+        >>> print(device)
+        cuda:1
     """ # noqa E501
     if(device_name != ''):
         device = torch.device(device_name)
diff --git a/lightwood/helpers/torch.py b/lightwood/helpers/torch.py
--- a/lightwood/helpers/torch.py
+++ b/lightwood/helpers/torch.py
@@ -5,6 +5,32 @@
 
 
 def concat_vectors_and_pad(vec_list, max_):
+    """
+    Concatenates a list of input vectors and pads them to match a specified maximum
+    length.
+
+    This function takes a list of input vectors, concatenates them along a specified
+    dimension (dim=0), and then pads the concatenated vector to achieve a specified
+    maximum length. The padding is done with zeros.
+
+    Args:
+        vec_list (list of torch.Tensor): List of input vectors to concatenate and pad.
+        max_ (int): The maximum length of the concatenated and padded vector.
+
+    Returns:
+        torch.Tensor: The concatenated and padded vector.
+
+    Raises:
+        AssertionError: If the length of 'vec_list' is not greater than 0, or if it
+                        exceeds 'max_len', or if 'max_len' is not greater than 0.
+
+    Example:
+        >>> input_tensors = [torch.tensor([1, 2]), torch.tensor([3, 4, 5])]
+        >>> max_length = 5
+        >>> concatenated_padded = concat_vectors_and_pad(input_tensors, max_length)
+        >>> print(concatenated_padded)
+        tensor([1, 2, 3, 4, 5])
+    """
     assert len(vec_list) > 0
     assert len(vec_list) <= max_
     assert max_ > 0
@@ -27,10 +53,29 @@ class LightwoodAutocast:
     """
     Equivalent to torch.cuda.amp.autocast, but checks device compute capability
     to activate the feature only when the GPU has tensor cores to leverage AMP.
+
+    **Attributes:**
+
+    * `active` (bool): Whether AMP is currently active. This attribute is at the class
+    level
+
+    **Usage:**
+
+    ```python
+    >>> import lightwood.helpers.torch as lt
+    >>> with lt.LightwoodAutocast():
+    ...     # This code will be executed in AMP mode.
+    ...     pass
     """
     active = False
 
     def __init__(self, enabled=True):
+        """
+        Initializes the context manager for Automatic Mixed Precision (AMP) functionality.
+
+        Args:
+            enabled (bool, optional): Whether to enable AMP. Defaults to True.
+        """
         self.major = 0  # GPU major version
         torch_version = [int(i) for i in torch.__version__.split('.')[:-1]]
 
@@ -50,12 +95,18 @@ def __init__(self, enabled=True):
         LightwoodAutocast.active = self._enabled
 
     def __enter__(self):
+        """
+        * `__enter__()`: Enters the context manager and enables AMP if it is not already enabled.
+        """
         if self._enabled:
             self.prev = torch.is_autocast_enabled()
             torch.set_autocast_enabled(self._enabled)
             torch.autocast_increment_nesting()
 
     def __exit__(self, *args):
+        """
+        * `__exit__()`: Exits the context manager and disables AMP.
+        """
         if self._enabled:
             # Drop the cache when we exit to a nesting level that's outside any instance of autocast
             if torch.autocast_decrement_nesting() == 0:
@@ -64,6 +115,9 @@ def __exit__(self, *args):
         return False
 
     def __call__(self, func):
+        """
+        * `__call__(self, func)`: Returns a decorated function that enables AMP when it is called.
+        """
         @functools.wraps(func)
         def decorate_autocast(*args, **kwargs):
             with self:


Test patch: