diff --git a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
index 01d7d73fc..29da449ff 100644
--- a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
+++ b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
@@ -3996,8 +3996,9 @@ export default function MyFlexUpdateForm(props) {
           modelFields = onSubmit(modelFields);
         }
         try {
+          const original = await DataStore.query(Flex0, id);
           await DataStore.save(
-            Flex0.copyOf(flexRecord, (updated) => {
+            Flex0.copyOf(original, (updated) => {
               Object.assign(updated, modelFields);
             })
           );
diff --git a/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts b/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
index 8a31472a8..2e705bf1e 100644
--- a/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
+++ b/packages/codegen-ui-react/lib/amplify-ui-renderers/form.ts
@@ -127,7 +127,7 @@ export default class FormRenderer extends ReactComponentRenderer<BaseComponentPr
         factory.createTryStatement(
           factory.createBlock(
             [
-              ...buildDataStoreExpression(formActionType, dataTypeName, importedModelName),
+              ...buildDataStoreExpression(formActionType, importedModelName),
               // call onSuccess hook if it exists
               factory.createIfStatement(
                 factory.createIdentifier('onSuccess'),
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
index 9e2f997ae..fda836245 100644
--- a/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper.ts
@@ -546,11 +546,7 @@ export const buildOnChangeStatement = (
   );
 };
 
-export const buildDataStoreExpression = (
-  dataStoreActionType: 'update' | 'create',
-  dataTypeName: string,
-  importedModelName: string,
-) => {
+export const buildDataStoreExpression = (dataStoreActionType: 'update' | 'create', importedModelName: string) => {
   if (dataStoreActionType === 'update') {
     return [
       factory.createVariableStatement(
@@ -568,7 +564,7 @@ export const buildDataStoreExpression = (
                     factory.createIdentifier('query'),
                   ),
                   undefined,
-                  [factory.createIdentifier(modelName), factory.createIdentifier('id')],
+                  [factory.createIdentifier(importedModelName), factory.createIdentifier('id')],
                 ),
               ),
             ),