diff --git a/pytoda/smiles/transforms.py b/pytoda/smiles/transforms.py
index 6a9288f..a40439f 100644
--- a/pytoda/smiles/transforms.py
+++ b/pytoda/smiles/transforms.py
@@ -95,7 +95,7 @@ class ToTensor(Transform):
         """
         return torch.tensor(
             token_indexes, dtype=self.dtype, device=self.device
-        ).view(-1, 1)
+        ).view(-1, 1).squeeze()
 
 
 class RemoveIsomery(Transform):

Shape of tensors returned by SMILES_dataset
The `SMILESDataset` class currently returns per default tensors of shape `sequence_length x 1`. This is not ideal since the additional dimension is not understood by most models per default.
As a workaround, we use `torch.squeeze()` in the code that currently uses `SMILESDataset ` objects.

Writing methods that want to use models and loaders generically (like the uncertainity estimation functions in `paccmann_predictor`) causes problems then.

Check `test_smiles_eager_dataset.py`  for the details. 
diff --git a/pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py b/pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py
index f3033fb..10861c0 100644
--- a/pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py
+++ b/pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py
@@ -111,10 +111,7 @@ class TestDrugSensitivityDataset(unittest.TestCase):
                     np.testing.assert_almost_equal(
                         token_indexes_tensor.numpy(),
                         np.array(
-                            [
-                                [padding_index], [padding_index], [c_index],
-                                [o_index]
-                            ]
+                            [padding_index, padding_index, c_index, o_index]
                         )
                     )
                     np.testing.assert_almost_equal(
@@ -176,7 +173,7 @@ class TestDrugSensitivityDataset(unittest.TestCase):
                             ic50_batch
                         )
                     ) in enumerate(data_loader):
-                        self.assertEqual(token_indexes_batch.size(), (2, 4, 1))
+                        self.assertEqual(token_indexes_batch.size(), (2, 4))
                         self.assertEqual(gene_expression_batch.size(), (2, 4))
                         self.assertEqual(ic50_batch.size(), (2, 1))
                         if batch_index > 4:
diff --git a/pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py b/pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py
index cbd7c76..8043634 100644
--- a/pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py
+++ b/pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py
@@ -111,10 +111,7 @@ class TestDrugSensitivityDataset(unittest.TestCase):
                     np.testing.assert_almost_equal(
                         token_indexes_tensor.numpy(),
                         np.array(
-                            [
-                                [padding_index], [padding_index], [c_index],
-                                [o_index]
-                            ]
+                            [padding_index, padding_index, c_index, o_index]
                         )
                     )
                     np.testing.assert_almost_equal(
@@ -176,7 +173,7 @@ class TestDrugSensitivityDataset(unittest.TestCase):
                             ic50_batch
                         )
                     ) in enumerate(data_loader):
-                        self.assertEqual(token_indexes_batch.size(), (2, 4, 1))
+                        self.assertEqual(token_indexes_batch.size(), (2, 4))
                         self.assertEqual(gene_expression_batch.size(), (2, 4))
                         self.assertEqual(ic50_batch.size(), (2, 1))
                         if batch_index > 4:
diff --git a/pytoda/datasets/tests/test_smiles_eager_dataset.py b/pytoda/datasets/tests/test_smiles_eager_dataset.py
index 2479c3c..7c79c5b 100644
--- a/pytoda/datasets/tests/test_smiles_eager_dataset.py
+++ b/pytoda/datasets/tests/test_smiles_eager_dataset.py
@@ -169,7 +169,7 @@ class TestSMILESDatasetEagerBackend(unittest.TestCase):
                     smiles_dataset, batch_size=4, shuffle=True
                 )
                 for batch_index, batch in enumerate(data_loader):
-                    self.assertEqual(batch.shape, (4, 4, 1))
+                    self.assertEqual(batch.shape, (4, 4))
                     if batch_index > 10:
                         break
 
diff --git a/pytoda/datasets/tests/test_smiles_lazy_dataset.py b/pytoda/datasets/tests/test_smiles_lazy_dataset.py
index 01b5929..d2791f3 100644
--- a/pytoda/datasets/tests/test_smiles_lazy_dataset.py
+++ b/pytoda/datasets/tests/test_smiles_lazy_dataset.py
@@ -169,7 +169,7 @@ class TestSMILESDatasetLazyBackend(unittest.TestCase):
                     smiles_dataset, batch_size=4, shuffle=True
                 )
                 for batch_index, batch in enumerate(data_loader):
-                    self.assertEqual(batch.shape, (4, 4, 1))
+                    self.assertEqual(batch.shape, (4, 4))
                     if batch_index > 10:
                         break
 

conda.yml
pip install -e .
parse_log_pytest
environment.yml
pytest
apt-get update apt-get install -y gcc
3.7
pytest --no-header -rA --tb=line --color=no -p no:cacheprovider -W ignore::DeprecationWarning
Bottleneck @ file:///opt/conda/conda-bld/bottleneck_1657175564434/work
certifi @ file:///croot/certifi_1671487769961/work/certifi
dill==0.3.1.1
diskcache==4.1.0
exceptiongroup==1.2.2
importlib-metadata==6.7.0
iniconfig==2.0.0
joblib==1.3.2
numexpr @ file:///croot/numexpr_1668713893690/work
numpy==1.16.1
packaging @ file:///croot/packaging_1671697413597/work
pandas==0.25.2
Pillow==9.4.0
pluggy==1.2.0
pytest==7.4.4
python-dateutil @ file:///tmp/build/80754af9/python-dateutil_1626374649649/work
-e git+https://github.com/PaccMann/paccmann_datasets.git@def0673e39f2e6a6e86f6d2456eab4e7b770980f#egg=pytoda
pytz @ file:///croot/pytz_1671697431263/work
scikit-learn==0.21.3
scipy==1.5.4
selfies==0.2.4
six @ file:///tmp/build/80754af9/six_1644875935023/work
tomli==2.0.1
torch==1.0.1
typing_extensions==4.7.1
zipp==3.15.0

name: paccmann_datasets
channels:
  - rdkit
  - defaults
  - https://repo.anaconda.com/pkgs/main
  - https://repo.anaconda.com/pkgs/r
  - conda-forge
dependencies:
  - _libgcc_mutex=0.1=main
  - _openmp_mutex=5.1=1_gnu
  - blas=1.0=openblas
  - bottleneck=1.3.5=py37h7deecbd_0
  - bzip2=1.0.8=h5eee18b_6
  - ca-certificates=2025.2.25=h06a4308_0
  - cairo=1.16.0=hb05425b_5
  - certifi=2022.12.7=py37h06a4308_0
  - expat=2.6.4=h6a678d5_0
  - fontconfig=2.14.1=h4c34cd2_2
  - freetype=2.12.1=h4a9f257_0
  - giflib=5.2.2=h5eee18b_0
  - glib=2.78.4=h6a678d5_0
  - glib-tools=2.78.4=h6a678d5_0
  - icu=58.2=he6710b0_3
  - jpeg=9e=h5eee18b_3
  - lcms2=2.16=hb9589c4_0
  - ld_impl_linux-64=2.40=h12ee557_0
  - lerc=4.0.0=h6a678d5_0
  - libboost=1.67.0=h46d08c1_4
  - libdeflate=1.22=h5eee18b_0
  - libffi=3.4.4=h6a678d5_1
  - libgcc-ng=11.2.0=h1234567_1
  - libgfortran-ng=11.2.0=h00389a5_1
  - libgfortran5=11.2.0=h1234567_1
  - libglib=2.78.4=hdc74915_0
  - libgomp=11.2.0=h1234567_1
  - libiconv=1.16=h5eee18b_3
  - libopenblas=0.3.21=h043d6bf_0
  - libpng=1.6.39=h5eee18b_0
  - libstdcxx-ng=11.2.0=h1234567_1
  - libtiff=4.5.1=hffd6297_1
  - libuuid=1.41.5=h5eee18b_0
  - libwebp=1.2.4=h11a3e52_1
  - libwebp-base=1.2.4=h5eee18b_1
  - libxcb=1.15=h7f8727e_0
  - libxml2=2.10.4=hcbfbd50_0
  - lz4-c=1.9.4=h6a678d5_1
  - ncurses=6.4=h6a678d5_0
  - numexpr=2.8.4=py37hd2a5715_0
  - openssl=1.1.1w=h7f8727e_0
  - packaging=22.0=py37h06a4308_0
  - pcre2=10.42=hebb0a14_1
  - pillow=9.4.0=py37h6a678d5_0
  - pip=22.3.1=py37h06a4308_0
  - pixman=0.40.0=h7f8727e_1
  - py-boost=1.67.0=py37h04863e7_4
  - python=3.7.16=h7a1cb2a_0
  - python-dateutil=2.8.2=pyhd3eb1b0_0
  - pytz=2022.7=py37h06a4308_0
  - rdkit=2019.03.1.0=py37hc20afe1_1
  - readline=8.2=h5eee18b_0
  - setuptools=65.6.3=py37h06a4308_0
  - six=1.16.0=pyhd3eb1b0_1
  - sqlite=3.45.3=h5eee18b_0
  - tk=8.6.14=h39e8969_0
  - wheel=0.38.4=py37h06a4308_0
  - xz=5.6.4=h5eee18b_1
  - zlib=1.2.13=h5eee18b_1
  - zstd=1.5.6=hc292b87_0
  - pip:
      - dill==0.3.1.1
      - diskcache==4.1.0
      - exceptiongroup==1.2.2
      - importlib-metadata==6.7.0
      - iniconfig==2.0.0
      - joblib==1.3.2
      - numpy==1.16.1
      - pandas==0.25.2
      - pluggy==1.2.0
      - pytest==7.4.4
      - scikit-learn==0.21.3
      - scipy==1.5.4
      - selfies==0.2.4
      - tomli==2.0.1
      - torch==1.0.1
      - typing-extensions==4.7.1
      - zipp==3.15.0
prefix: /opt/conda/envs/paccmann_datasets

pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py::TestDrugSensitivityDataset::test___getitem__ pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py::TestDrugSensitivityDataset::test_data_loader pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py::TestDrugSensitivityDataset::test___getitem__ pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py::TestDrugSensitivityDataset::test_data_loader pytoda/datasets/tests/test_smiles_eager_dataset.py::TestSMILESDatasetEagerBackend::test_data_loader pytoda/datasets/tests/test_smiles_lazy_dataset.py::TestSMILESDatasetLazyBackend::test_data_loader
pytoda/datasets/tests/test_drug_sensitivity_eager_dataset.py::TestDrugSensitivityDataset::test___len__ pytoda/datasets/tests/test_drug_sensitivity_lazy_dataset.py::TestDrugSensitivityDataset::test___len__ pytoda/datasets/tests/test_smiles_eager_dataset.py::TestSMILESDatasetEagerBackend::test___getitem__ pytoda/datasets/tests/test_smiles_eager_dataset.py::TestSMILESDatasetEagerBackend::test___len__ pytoda/datasets/tests/test_smiles_lazy_dataset.py::TestSMILESDatasetLazyBackend::test___getitem__ pytoda/datasets/tests/test_smiles_lazy_dataset.py::TestSMILESDatasetLazyBackend::test___len__
pytoda/smiles/transforms.py
pytoda/smiles/transforms.py