diff --git a/src/input/dataSchema.ts b/src/input/dataSchema.ts
index 7d033ebf..5cfd3515 100644
--- a/src/input/dataSchema.ts
+++ b/src/input/dataSchema.ts
@@ -126,8 +126,8 @@ export class DataSchema {
   constructor(dataSchema: StringDict | string) {
     if (typeof dataSchema === "string") {
       this.replace = new DataSchemaReplace(JSON.parse(dataSchema)["replace"]);
-    } else if (dataSchema["replace"] instanceof DataSchemaReplace) {
-      this.replace = dataSchema["replace"];
+    } else if (dataSchema instanceof DataSchema) {
+      this.replace = dataSchema.replace;
     } else {
       this.replace = new DataSchemaReplace(dataSchema["replace"] as StringDict);
     }
diff --git a/tests/v2/input/inferenceParameter.spec.ts b/tests/v2/input/inferenceParameter.spec.ts
index cad3056c..51a48171 100644
--- a/tests/v2/input/inferenceParameter.spec.ts
+++ b/tests/v2/input/inferenceParameter.spec.ts
@@ -18,31 +18,32 @@ describe("MindeeV2 - Inference Parameter", () => {
     expectedDataSchemaObject = new DataSchema(expectedDataSchemaDict);
   });
 
-  it("shouldn't replace when unset", async () => {
-    const params: InferenceParameters = {
-      modelId: "test-model-id",
-    };
+  describe("dataSchema", () => {
+    it("shouldn't replace when unset", async () => {
+      const params: InferenceParameters = {
+        modelId: "test-model-id",
+      };
 
-    expect(params.dataSchema).to.be.undefined;
-  });
+      expect(params.dataSchema).to.be.undefined;
+    });
 
-  it("should equate no matter the type", async () => {
-    const paramsDict: InferenceParameters = {
-      modelId: "test-model-id",
-      dataSchema: expectedDataSchemaDict,
-    };
-    const paramsString: InferenceParameters = {
-      modelId: "test-model-id",
-      dataSchema: expectedDataSchemaString,
-    };
-    const paramsObject: InferenceParameters = {
-      modelId: "test-model-id",
-      dataSchema: expectedDataSchemaObject,
-    };
+    it("should equate no matter the type", async () => {
+      const paramsDict: InferenceParameters = {
+        modelId: "test-model-id",
+        dataSchema: expectedDataSchemaDict,
+      };
+      const paramsString: InferenceParameters = {
+        modelId: "test-model-id",
+        dataSchema: expectedDataSchemaString,
+      };
+      const paramsObject: InferenceParameters = {
+        modelId: "test-model-id",
+        dataSchema: expectedDataSchemaObject,
+      };
 
-    expect(JSON.stringify(paramsDict.dataSchema)).to.eq(expectedDataSchemaString);
-    expect(paramsObject.dataSchema?.toString()).to.eq(expectedDataSchemaString);
-    expect(paramsString.dataSchema?.toString()).to.eq(expectedDataSchemaString);
+      expect(JSON.stringify(paramsDict.dataSchema)).to.eq(expectedDataSchemaString);
+      expect(paramsObject.dataSchema?.toString()).to.eq(expectedDataSchemaString);
+      expect(paramsString.dataSchema?.toString()).to.eq(expectedDataSchemaString);
+    });
   });
-
 });