diff --git a/tests/data/diffusion/test_data_loader.py b/tests/data/diffusion/test_data_loader.py
index bc8d7537..a811dcd1 100644
--- a/tests/data/diffusion/test_data_loader.py
+++ b/tests/data/diffusion/test_data_loader.py
@@ -1,7 +1,7 @@
 import pytest
 import torch
 
-from crystal_diffusion.data.diffusion.data_loader import LammpsForDiffusionDataModule  # Replace 'your_module' with the actual module name
+from crystal_diffusion.data.diffusion.data_loader import LammpsForDiffusionDataModule
 
 
 @pytest.fixture