diff --git a/SIRF_data_preparation/create_initial_images.py b/SIRF_data_preparation/create_initial_images.py
index f8b13a2..5dc5ad4 100644
--- a/SIRF_data_preparation/create_initial_images.py
+++ b/SIRF_data_preparation/create_initial_images.py
@@ -51,7 +51,10 @@ def scale_initial_image(acquired_data, additive_term, mult_factors, template_ima
 
     WARNING: assumes that obj_fun has been set_up already
     """
-    data_sum = (acquired_data.sum() - (additive_term * mult_factors).sum())
+    sensitivity_factors = STIR.AcquisitionSensitivityModel(mult_factors)
+    sensitivity_factors.set_up(acquired_data)
+    background = sensitivity_factors.forward(additive_term)
+    data_sum = acquired_data.sum() - background.sum()
     if data_sum <= 0 or math.isinf(data_sum) or math.isnan(data_sum):
         raise ValueError("Something wrong with input data. Sum of (prompts-background) is negative:"
                          f" sum prompts: {acquired_data.sum()}, sum corrected: {data_sum}")
diff --git a/main_OSEM.py b/main_OSEM.py
index c68e72d..b16ecd2 100644
--- a/main_OSEM.py
+++ b/main_OSEM.py
@@ -63,9 +63,16 @@ def __init__(self, data: Dataset, num_subsets: int = 7, update_objective_interva
 
             acquisition_model_subset = STIR.AcquisitionModelUsingParallelproj()
             acquisition_model_subset.set_additive_term(additive_term_subset)
-            acquisition_model_subset.set_up(prompts_subset, self.x)
+            if prompts_subset.dimensions()[0] > 1 and multiplicative_factors_subset.dimensions()[0] == 1:
+                # non-TOF multiplicative factors for TOF data
+                acquisition_model_subset.set_up(multiplicative_factors_subset, self.x)
+                subset_sensitivity = acquisition_model_subset.backward(multiplicative_factors_subset)
+                acquisition_model_subset.set_up(prompts_subset, self.x)
+            else:
+                assert prompts_subset.dimensions()[0] == multiplicative_factors_subset.dimensions()[0]
+                acquisition_model_subset.set_up(prompts_subset, self.x)
+                subset_sensitivity = acquisition_model_subset.backward(multiplicative_factors_subset)
 
-            subset_sensitivity = acquisition_model_subset.backward(multiplicative_factors_subset)
             # add a small number to avoid NaN in division
             subset_sensitivity += subset_sensitivity.max() * 1e-6