diff --git a/model.py b/model.py
index c246fda..b0603b2 100644
--- a/model.py
+++ b/model.py
@@ -13,11 +13,13 @@ class HeatPINN(nn.Module):
         layers.append(nn.Linear(h, 1))
         self.net = nn.Sequential(*layers)
 
-    def forward(self, x):
+    def forward(self, xt):
+        # Normalize inputs to [0,1]x[0,1] so both dimensions share the same scale
+        scale = torch.tensor([config.L, config.T_END], device=xt.device, dtype=xt.dtype)
+        xt_norm = xt / scale
         # Output scaled to physical range: T_AMB + (Q*L/K) * net
         # net learns dimensionless perturbation in [0,1] range
-        T_scale = config.T_AMB + (config.Q_VAL * config.L / config.K) * self.net(x)
-        return T_scale
+        return config.T_AMB + (config.Q_VAL * config.L / config.K) * self.net(xt_norm)
 
 
 def heat_pinn_loss(model, x_f, t_f, x_ic, t_bc,
@@ -33,8 +35,7 @@ def heat_pinn_loss(model, x_f, t_f, x_ic, t_bc,
     x_f = x_f.detach().requires_grad_(True)
     t_f = t_f.detach().requires_grad_(True)
     T_f = model(torch.stack([x_f, t_f], dim=1))
-    dT_dt = torch.autograd.grad(T_f.sum(), t_f, create_graph=True)[0]
-    dT_dx = torch.autograd.grad(T_f.sum(), x_f, create_graph=True)[0]
+    dT_dt, dT_dx = torch.autograd.grad(T_f.sum(), [t_f, x_f], create_graph=True)
     d2T_dx2 = torch.autograd.grad(dT_dx.sum(), x_f, create_graph=True)[0]
     sigma = 0.02
     Q_t_f = torch.where(t_f >= config.T_STEP,