diff --git a/Makefile b/Makefile
index 4024f500a87b1a901a76683ea46b9194ee518ded..1a0887c774707bcf82dc013a562d98daf26905f1 100644
--- a/Makefile
+++ b/Makefile
@@ -17,11 +17,9 @@ DEBUG_FLAGS=
 # CUDA settings
 NVCC=nvcc
 NVCC_FLAGS=-O3 --use_fast_math
-#NVCC_FLAGS=-ccbin mpicc $(CFLAGS)
 NVCC_PATH:="$(shell which ${NVCC})"
 CUDA_FLAGS=-DENABLE_CUDA_AWARE_MPI
-CUDA_BIN_PATH:="$(shell dirname ${NVCC_PATH})"
-CUDA_PATH:="$(shell dirname ${CUDA_BIN_PATH})"
+CUDART_FLAGS=-lcudart -L /apps/SPACK/0.19.1/opt/linux-almalinux8-zen/gcc-8.5.0/nvhpc-23.7-bzxcokzjvx4stynglo4u2ffpljajzlam/Linux_x86_64/23.7/cuda/12.2/targets/x86_64-linux/lib
 
 # MPI settings
 MPI_PATH=/apps/SPACK/0.19.1/opt/linux-almalinux8-zen/intel-2021.10.0/openmpi-4.1.6-ijsnjhq77rjc256wlrp52m37rsq6miff
@@ -82,7 +80,7 @@ $(CPU_BIN): $(CPU_SRC) $(CPU_OBJ_PATH)/pairs.o $(CPU_OBJ_PATH)/regular_6d_stenci
 
 $(GPU_BIN): $(GPU_SRC) $(GPU_OBJ_PATH)/pairs.o $(GPU_OBJ_PATH)/regular_6d_stencil.o $(GPU_OBJ_PATH)/cuda_runtime.o
 	$(NVCC) $(NVCC_FLAGS) -c -o $(GPU_OBJ_PATH)/$(GPU_BIN).o $(GPU_SRC) $(DEBUG_FLAGS) $(MPI_FLAGS) $(CUDA_FLAGS)
-	$(CC) -o $(GPU_BIN) $(GPU_OBJ_PATH)/$(GPU_BIN).o $(GPU_OBJ_PATH)/cuda_runtime.o $(GPU_OBJ_PATH)/pairs.o $(GPU_OBJ_PATH)/regular_6d_stencil.o -lcudart -L$(CUDA_PATH)/lib64 $(CUDA_FLAGS) $(CFLAGS)
+	$(CC) -o $(GPU_BIN) $(GPU_OBJ_PATH)/$(GPU_BIN).o $(GPU_OBJ_PATH)/cuda_runtime.o $(GPU_OBJ_PATH)/pairs.o $(GPU_OBJ_PATH)/regular_6d_stencil.o $(CUDART_FLAGS) $(CUDA_FLAGS) $(CFLAGS)
 
 clean:
 	@echo "Cleaning..."
diff --git a/examples/dem.py b/examples/dem.py
index 076660c79fc4eef8d4f97ce2eb4fe6a17862cee0..0276de9b5bd338b6f05795026a65b749a345d308 100644
--- a/examples/dem.py
+++ b/examples/dem.py
@@ -97,8 +97,9 @@ if target != 'cpu' and target != 'gpu':
     print(f"Invalid target, use {cmd} <cpu/gpu>")
 
 # Config file parameters
-domainSize_SI = [0.8, 0.015, 0.2]
-blocks = [3, 3, 1]
+#domainSize_SI = [0.8, 0.015, 0.2]
+domainSize_SI = [0.6, 0.6, 0.2] # node base
+#domainSize_SI = [0.8, 0.8, 0.2] # node base
 diameter_SI = 0.0029
 gravity_SI = 9.81
 densityFluid_SI = 1000
@@ -111,6 +112,7 @@ restitutionCoefficient = 0.1
 collisionTime_SI = 5e-4
 poissonsRatio = 0.22
 timeSteps = 10000
+#timeSteps = 1000
 visSpacing = 100
 denseBottomLayer = False
 bottomLayerOffsetFactor = 1.0
@@ -130,7 +132,9 @@ psim = pairs.simulation(
     [pairs.sphere(), pairs.halfspace()],
     timesteps=timeSteps,
     double_prec=True,
-    use_contact_history=True)
+    use_contact_history=True,
+    particle_capacity=1000000,
+    neighbor_capacity=20)
 
 if target == 'gpu':
     psim.target(pairs.target_gpu())
@@ -159,7 +163,10 @@ psim.set_domain([0.0, 0.0, 0.0, domainSize_SI[0], domainSize_SI[1], domainSize_S
 psim.set_domain_partitioner(pairs.regular_domain_partitioner_xy())
 psim.pbc([True, True, False])
 psim.read_particle_data(
-    "data/spheres.input",
+    #"data/spheres.input",
+    #"data/spheres_4x4x2.input",
+    "data/spheres_6x6x2.input",
+    #"data/spheres_8x8x2.input",
     ['uid', 'type', 'mass', 'radius', 'position', 'linear_velocity', 'flags'],
     pairs.sphere())
 
diff --git a/src/pairs/__init__.py b/src/pairs/__init__.py
index 01126402422b66000359e9fffc7f964288d3df73..ab81066421f8c7f359c2a35694009b229395ef6d 100644
--- a/src/pairs/__init__.py
+++ b/src/pairs/__init__.py
@@ -6,8 +6,20 @@ from pairs.sim.shapes import Shapes
 from pairs.sim.simulation import Simulation
 
 
-def simulation(ref, shapes, dims=3, timesteps=100, double_prec=False, use_contact_history=False, debug=False):
-    return Simulation(CGen(ref, debug), shapes, dims, timesteps, double_prec, use_contact_history)
+def simulation(
+    ref,
+    shapes,
+    dims=3,
+    timesteps=100,
+    double_prec=False,
+    use_contact_history=False,
+    particle_capacity=800000,
+    neighbor_capacity=100,
+    debug=False):
+
+    return Simulation(
+        CGen(ref, debug), shapes, dims, timesteps, double_prec, use_contact_history,
+        particle_capacity, neighbor_capacity)
 
 def target_cpu():
     return Target(Target.Backend_CPP, Target.Feature_CPU)
diff --git a/src/pairs/sim/simulation.py b/src/pairs/sim/simulation.py
index a1fefa45c053670b66665e4169212ad0f71c0b8a..5bb602ccacfcb0877d922ab8e0cb06851886e44e 100644
--- a/src/pairs/sim/simulation.py
+++ b/src/pairs/sim/simulation.py
@@ -34,7 +34,17 @@ from pairs.transformations import Transformations
 
 
 class Simulation:
-    def __init__(self, code_gen, shapes, dims=3, timesteps=100, double_prec=False, use_contact_history=False):
+    def __init__(
+        self,
+        code_gen,
+        shapes,
+        dims=3,
+        timesteps=100,
+        double_prec=False,
+        use_contact_history=False,
+        particle_capacity=800000,
+        neighbor_capacity=100):
+
         self.code_gen = code_gen
         self.code_gen.assign_simulation(self)
         self.position_prop = None
@@ -44,8 +54,8 @@ class Simulation:
         self.features = Features(self)
         self.feature_properties = FeatureProperties(self)
         self.contact_properties = ContactProperties(self)
-        self.particle_capacity = self.add_var('particle_capacity', Types.Int32, 800000)
-        self.neighbor_capacity = self.add_var('neighbor_capacity', Types.Int32, 100)
+        self.particle_capacity = self.add_var('particle_capacity', Types.Int32, particle_capacity)
+        self.neighbor_capacity = self.add_var('neighbor_capacity', Types.Int32, neighbor_capacity)
         self.nlocal = self.add_var('nlocal', Types.Int32)
         self.nghost = self.add_var('nghost', Types.Int32)
         self.resizes = self.add_array('resizes', 3, Types.Int32, arr_sync=False)