diff --git a/src/products/components/ProductCreatePage/ProductCreatePage.tsx b/src/products/components/ProductCreatePage/ProductCreatePage.tsx
index 4c30f8fb6..ed55e6f6d 100644
--- a/src/products/components/ProductCreatePage/ProductCreatePage.tsx
+++ b/src/products/components/ProductCreatePage/ProductCreatePage.tsx
@@ -35,7 +35,10 @@ import useStateFromProps from "@saleor/hooks/useStateFromProps";
 import { sectionNames } from "@saleor/intl";
 import { ConfirmButtonTransitionState } from "@saleor/macaw-ui";
 import ProductVariantPrice from "@saleor/products/components/ProductVariantPrice";
-import { productListUrl } from "@saleor/products/urls";
+import {
+  ProductCreateUrlQueryParams,
+  productListUrl,
+} from "@saleor/products/urls";
 import { getChoices } from "@saleor/products/utils/data";
 import React from "react";
 import { useIntl } from "react-intl";
@@ -92,7 +95,7 @@ interface ProductCreatePageProps {
   fetchMoreReferencePages?: FetchMoreProps;
   fetchMoreReferenceProducts?: FetchMoreProps;
   onAttributeSelectBlur: () => void;
-  onCloseDialog: () => void;
+  onCloseDialog: (currentParams?: ProductCreateUrlQueryParams) => void;
   onSelectProductType: (productTypeId: string) => void;
   onSubmit?(data: ProductCreateData);
 }
@@ -142,6 +145,10 @@ export const ProductCreatePage: React.FC<ProductCreatePageProps> = ({
   const intl = useIntl();
   const navigate = useNavigator();
 
+  const closeDialog = () => {
+    onCloseDialog({ "product-type-id": selectedProductType.id });
+  };
+
   // Display values
   const [selectedCategory, setSelectedCategory] = useStateFromProps(
     initial?.category || "",
@@ -164,7 +171,8 @@ export const ProductCreatePage: React.FC<ProductCreatePageProps> = ({
       value: taxClass.id,
     })) ?? [];
 
-  const canOpenAssignReferencesAttributeDialog = !!assignReferencesAttributeId;
+  const canOpenAssignReferencesAttributeDialog =
+    !!assignReferencesAttributeId && !fetchMoreReferenceProducts.loading;
 
   const handleAssignReferenceAttribute = (
     attributeValues: string[],
@@ -179,7 +187,8 @@ export const ProductCreatePage: React.FC<ProductCreatePageProps> = ({
         data.attributes,
       ),
     );
-    onCloseDialog();
+
+    closeDialog();
   };
 
   return (
@@ -222,7 +231,6 @@ export const ProductCreatePage: React.FC<ProductCreatePageProps> = ({
         const isSimpleProduct = data.productType?.hasVariants === false;
 
         const errors = [...apiErrors, ...validationErrors];
-
         return (
           <Container>
             <Backlink href={productListUrl()}>
@@ -393,7 +401,7 @@ export const ProductCreatePage: React.FC<ProductCreatePageProps> = ({
                 onFetch={handlers.fetchReferences}
                 onFetchMore={handlers.fetchMoreReferences?.onFetchMore}
                 loading={handlers.fetchMoreReferences?.loading}
-                onClose={onCloseDialog}
+                onClose={closeDialog}
                 onSubmit={attributeValues =>
                   handleAssignReferenceAttribute(
                     attributeValues,
diff --git a/src/products/components/ProductUpdatePage/ProductUpdatePage.tsx b/src/products/components/ProductUpdatePage/ProductUpdatePage.tsx
index 9fd4e1970..89709341d 100644
--- a/src/products/components/ProductUpdatePage/ProductUpdatePage.tsx
+++ b/src/products/components/ProductUpdatePage/ProductUpdatePage.tsx
@@ -199,7 +199,8 @@ export const ProductUpdatePage: React.FC<ProductUpdatePageProps> = ({
       value: taxClass.id,
     })) || [];
 
-  const canOpenAssignReferencesAttributeDialog = !!assignReferencesAttributeId;
+  const canOpenAssignReferencesAttributeDialog =
+    !!assignReferencesAttributeId && !fetchMoreReferenceProducts.loading;
 
   const handleAssignReferenceAttribute = (
     attributeValues: string[],
diff --git a/src/products/views/ProductCreate/ProductCreate.tsx b/src/products/views/ProductCreate/ProductCreate.tsx
index e60fc59e8..881a7b821 100644
--- a/src/products/views/ProductCreate/ProductCreate.tsx
+++ b/src/products/views/ProductCreate/ProductCreate.tsx
@@ -239,6 +239,7 @@ export const ProductCreateView: React.FC<ProductCreateProps> = ({ params }) => {
   const handleAssignAttributeReferenceClick = (attribute: AttributeInput) =>
     navigate(
       productAddUrl({
+        ...params,
         action: "assign-attribute-value",
         id: attribute.id,
       }),
@@ -373,7 +374,7 @@ export const ProductCreateView: React.FC<ProductCreateProps> = ({ params }) => {
         fetchReferenceProducts={searchProducts}
         fetchMoreReferenceProducts={fetchMoreReferenceProducts}
         fetchMoreAttributeValues={fetchMoreAttributeValues}
-        onCloseDialog={() => navigate(productAddUrl())}
+        onCloseDialog={currentParams => navigate(productAddUrl(currentParams))}
         selectedProductType={selectedProductType?.productType}
         onSelectProductType={handleSelectProductType}
         onAttributeSelectBlur={searchAttributeReset}