diff --git a/litex_boards/partner/targets/ulx3s.py b/litex_boards/partner/targets/ulx3s.py
index 4e2e8b1..90bc5e9 100755
--- a/litex_boards/partner/targets/ulx3s.py
+++ b/litex_boards/partner/targets/ulx3s.py
@@ -16,7 +16,7 @@ from litex.soc.cores.clock import *
 from litex.soc.integration.soc_sdram import *
 from litex.soc.integration.builder import *
 
-from litedram.modules import MT48LC16M16, AS4C32M16, AS4C16M16
+from litedram import modules as litedram_modules
 from litedram.phy import GENSDRPHY
 
 # CRG ----------------------------------------------------------------------------------------------
@@ -55,7 +55,7 @@ class _CRG(Module):
 
 class BaseSoC(SoCSDRAM):
     def __init__(self, device="LFE5U-45F", toolchain="diamond",
-        sys_clk_freq=int(50e6), mem_device="MT48LC16M16", **kwargs):
+        sys_clk_freq=int(50e6), sdram_module_cls="MT48LC16M16", **kwargs):
 
         platform = ulx3s.Platform(device=device, toolchain=toolchain)
         SoCSDRAM.__init__(self, platform, clk_freq=sys_clk_freq,
@@ -65,13 +65,8 @@ class BaseSoC(SoCSDRAM):
         self.submodules.crg = _CRG(platform, sys_clk_freq)
 
         if not self.integrated_main_ram_size:
-            if mem_device.strip().upper() == "MT48LC16M16":
-                self.submodules.sdrphy = GENSDRPHY(platform.request("sdram"), cl=2)
-            else:
-                self.submodules.sdrphy = GENSDRPHY(platform.request("sdram"))
-            
-            memcls = getattr(sys.modules[__name__], mem_device.strip().upper())
-            sdram_module = memcls(sys_clk_freq, "1:1")
+            self.submodules.sdrphy = GENSDRPHY(platform.request("sdram"), cl=2)
+            sdram_module = getattr(litedram_modules, sdram_module_cls)(sys_clk_freq, "1:1")
             self.register_sdram(self.sdrphy,
                                 sdram_module.geom_settings,
                                 sdram_module.timing_settings)
@@ -86,15 +81,15 @@ def main():
         help='FPGA device, ULX3S can be populated with LFE5U-45F (default) or LFE5U-85F')
     parser.add_argument("--sys-clk-freq", default=50e6,
                         help="system clock frequency (default=50MHz)")
-    parser.add_argument("--mem-device", default="MT48LC16M16",
-                        help="Part number for SDRAM (default=MT48LC16M16)")
+    parser.add_argument("--sdram-module", default="MT48LC16M16",
+                        help="SDRAM module: MT48LC16M16, AS4C32M16 or AS4C16M16 (default=MT48LC16M16)")
     builder_args(parser)
     soc_sdram_args(parser)
     args = parser.parse_args()
 
     soc = BaseSoC(device=args.device, toolchain=args.toolchain,
         sys_clk_freq=int(float(args.sys_clk_freq)),
-        mem_device=args.mem_device,
+        sdram_module_cls=args.sdram_module,
         **soc_sdram_argdict(args))
     builder = Builder(soc, **builder_argdict(args))
     builder.build()