diff --git a/src/Template/AttributesSet.php b/src/Template/AttributesSet.php
index 39c7eb275815a29bff6f2042e836f3b832a4af48..8a05235e1958a7e7722d35fa3b649f54250106ff 100644
--- a/src/Template/AttributesSet.php
+++ b/src/Template/AttributesSet.php
@@ -45,9 +45,15 @@ class AttributesSet {
     }
     foreach ($style_set as $style_set_id => $value) {
       if (isset($this->attributesSetConfig[$style_set_id][$value])) {
-        foreach ($this->attributesSetConfig[$style_set_id][$value] as $style_key => $style_value) {
-          $classes_map[$style_key] = $style_value;
+        if (is_array($this->attributesSetConfig[$style_set_id][$value])) {
+          foreach ($this->attributesSetConfig[$style_set_id][$value] as $style_key => $style_value) {
+            $classes_map[$style_key] = $style_value;
+          }
         }
+        else {
+          $classes_map[$style_set_id] = $this->attributesSetConfig[$style_set_id][$value];
+        }
+
       }
     }
     return $classes_map;
diff --git a/tests/fixtures/AttributesSetDataSet.yml b/tests/fixtures/AttributesSetDataSet.yml
index 72240abf0ac999fe9be64e82cc80979d217fd36a..3d9a8162c8c07433d99d9e9d05f1788b0228bebc 100644
--- a/tests/fixtures/AttributesSetDataSet.yml
+++ b/tests/fixtures/AttributesSetDataSet.yml
@@ -4,6 +4,7 @@ style:
     background: ''
     text: ''
     border: ''
+    border_size: ''
   variant:
     primary:
       background: bg-blue
@@ -16,23 +17,30 @@ style:
       text: text-sm
     medium:
       text: text-lm
+  border_size:
+    small: border-sm
+    medium: border-md
 expected:
   -
     input:
       variant: primary
       button_size: small
+      border_size: small
     output:
       text: text-sm
       background: bg-blue
       border: bg-none
+      border_size: border-sm
   -
     input:
       variant: primary
       button_size: medium
+      border_size: small
     output:
       text: text-lm
       background: bg-blue
       border: bg-none
+      border_size: border-sm
   -
     input:
       variant: outline
@@ -41,14 +49,17 @@ expected:
       text: text-sm
       background: bg-none
       border: bg-gray
+      border_size: ''
   -
     input:
       variant: outline
+      border_size: medium
       button_size: medium
     output:
       text: text-lm
       background: bg-none
       border: bg-gray
+      border_size: border-md
   -
     input:
       variant: outline
@@ -56,3 +67,4 @@ expected:
       text: ''
       background: bg-none
       border: bg-gray
+      border_size: ''