Add tests and update documentation for devices field

Co-authored-by: manujosephv <[email protected]>

Author: Copilot

src/pytorch_tabular/config/config.py

@@ -287,8 +287,8 @@ class TrainerConfig:
                 'cpu','gpu','tpu','ipu', 'mps', 'auto'. Defaults to 'auto'.
                 Choices are: [`cpu`,`gpu`,`tpu`,`ipu`,'mps',`auto`].
 
-        devices (Optional[int]): Number of devices to train on (int). -1 uses all available devices. By
-                default, uses all available devices (-1)
+        devices (Union[int, List[int]]): Number of devices to train on (int), or list of device indices.
+                -1 uses all available devices. By default, uses all available devices (-1)
 
         devices_list (Optional[List[int]]): List of devices to train on (list). If specified, takes
                 precedence over `devices` argument. Defaults to None

tests/test_config.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python
+"""Tests for config classes."""
+
+import pytest
+from omegaconf import OmegaConf
+
+from pytorch_tabular.config import TrainerConfig
+
+
+class TestTrainerConfig:
+    """Tests for TrainerConfig class."""
+
+    def test_devices_list_to_devices_conversion(self):
+        """Test that devices_list is properly converted to devices."""
+        # Test with a list of devices
+        trainer_config = TrainerConfig(devices_list=[0, 1])
+        assert trainer_config.devices == [0, 1]
+        
+        # Wrap with OmegaConf as done in TabularModel
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == [0, 1]
+
+    def test_devices_list_multiple_gpus(self):
+        """Test devices_list with multiple GPU IDs as documented."""
+        trainer_config = TrainerConfig(devices_list=[1, 2, 3, 4])
+        assert trainer_config.devices == [1, 2, 3, 4]
+        
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == [1, 2, 3, 4]
+
+    def test_devices_int_value(self):
+        """Test that devices accepts integer values."""
+        trainer_config = TrainerConfig(devices=2)
+        assert trainer_config.devices == 2
+        
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == 2
+
+    def test_devices_default_value(self):
+        """Test that devices has default value of -1."""
+        trainer_config = TrainerConfig()
+        assert trainer_config.devices == -1
+        
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == -1
+
+    def test_devices_list_single_device(self):
+        """Test devices_list with a single device."""
+        trainer_config = TrainerConfig(devices_list=[0])
+        assert trainer_config.devices == [0]
+        
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == [0]
+
+    def test_devices_list_precedence(self):
+        """Test that devices_list takes precedence over devices."""
+        # When both are provided, devices_list should take precedence
+        trainer_config = TrainerConfig(devices=2, devices_list=[0, 1])
+        assert trainer_config.devices == [0, 1]
+        
+        config = OmegaConf.structured(trainer_config)
+        assert config.devices == [0, 1]
+
+    def test_omegaconf_merge_compatibility(self):
+        """Test that config works correctly with OmegaConf.merge."""
+        trainer_config = TrainerConfig(devices_list=[0, 1], max_epochs=10)
+        config = OmegaConf.structured(trainer_config)
+        
+        # Simulate merging as done in TabularModel
+        merged = OmegaConf.merge(
+            OmegaConf.to_container(config),
+            {"accelerator": "gpu"}
+        )
+        
+        assert merged.devices == [0, 1]
+        assert merged.max_epochs == 10
+        assert merged.accelerator == "gpu"