diff --git a/india_forecast_app/app.py b/india_forecast_app/app.py
index ade3297..be7ad87 100644
--- a/india_forecast_app/app.py
+++ b/india_forecast_app/app.py
@@ -84,12 +84,12 @@ def get_generation_data(
     generation_data = get_pv_generation_by_sites(
         session=db_session, site_uuids=site_uuids, start_utc=start, end_utc=end
     )
-    # hard code as for the moment
-    system_id = int(0.0)
+    # get the ml id, this only works for one site right now
+    system_id = sites[0].ml_id
 
     if len(generation_data) == 0:
         log.warning("No generation found for the specified sites/period")
-        generation_df = pd.DataFrame()
+        generation_df = pd.DataFrame(columns=[str(system_id)])
 
     else:
         # Convert to dataframe
@@ -272,7 +272,8 @@ def app_click(timestamp: dt.datetime | None, write_to_db: bool, log_level: str):
     Main click function for running forecasts for sites in India
     """
 
-    app(timestamp, write_to_db, log_level)
+    app(timestamp=timestamp, write_to_db=write_to_db, log_level=log_level)
+
 
 def app(timestamp: dt.datetime | None, write_to_db: bool=False, log_level: str="info"):
     """
diff --git a/india_forecast_app/models/pvnet/model.py b/india_forecast_app/models/pvnet/model.py
index 5facda0..8f5f538 100644
--- a/india_forecast_app/models/pvnet/model.py
+++ b/india_forecast_app/models/pvnet/model.py
@@ -61,7 +61,7 @@ def __init__(
         generation_data: dict[str, pd.DataFrame],
         hf_repo: str,
         hf_version: str,
-        name: str
+        name: str,
     ):
         """Initializer for the model"""
 
@@ -131,6 +131,7 @@ def predict(self, site_id: str, timestamp: dt.datetime):
             log.info("Feathering the forecast to the lastest value of generation")
 
             # Feather in the last generation, if it exists
+            system_id = str(self.generation_data["data"].columns[0])
             generation_da = self.generation_data["data"].to_xarray()
 
             # Check if the generation exists, if so, take the value at t0 and
@@ -139,7 +140,8 @@ def predict(self, site_id: str, timestamp: dt.datetime):
                 final_gen_points = 0
                 final_gen_index = 0
                 for gen_idx in range(len(generation_da.index.values) - 1, -1, -1):
-                    current_gen = generation_da.isel(index=gen_idx)["0"].values
+                    current_gen = generation_da.isel(index=gen_idx)
+                    current_gen = current_gen[system_id].values
                     if not np.isnan(current_gen) and current_gen > 0:
                         final_gen_points = current_gen * 1000.0
                         # Convert to KW back from MW
@@ -241,7 +243,10 @@ def _prepare_data_sources(self):
 
             # if generation_da is still empty make nans
             if len(generation_da) == 0:
-                generation_df = pd.DataFrame(index=forecast_timesteps, columns=["0"], data=0.0001)
+                cols = [str(col) for col in self.generation_data["data"].columns]
+                generation_df = pd.DataFrame(
+                    index=forecast_timesteps, columns=cols, data=0.0001
+                )
                 generation_da = generation_df.to_xarray()
             generation_da.to_netcdf(wind_netcdf_path, engine="h5netcdf")
 
@@ -271,7 +276,8 @@ def _prepare_data_sources(self):
 
             # if generation_da is still empty make nans
             if len(generation_da) == 0:
-                generation_df = pd.DataFrame(index=forecast_timesteps, columns=["0"], data=0.0001)
+                cols = [str(col) for col in self.generation_data["data"].columns]
+                generation_df = pd.DataFrame(index=forecast_timesteps, columns=cols, data=0.0001)
                 generation_da = generation_df.to_xarray()
             generation_da.to_netcdf(pv_netcdf_path, engine="h5netcdf")