Update model.py

model.py

@@ -55,4 +55,4 @@ class MedicalMasterAI(nn.Module):
         b, t = idx.shape
         x = self.token_embedding(idx) + self.position_embedding[:, :t, :]
         x = self.blocks(x)
-        return self.lm_head(self.ln_f(x)
+        return self.lm_head(self.ln_f(x))