diff --git a/builder/azure/arm/template_factory_test.TestVirtualMachineDeployment08.approved.json b/builder/azure/arm/template_factory_test.TestVirtualMachineDeployment08.approved.json
index 4370b37f1..8fa181797 100644
--- a/builder/azure/arm/template_factory_test.TestVirtualMachineDeployment08.approved.json
+++ b/builder/azure/arm/template_factory_test.TestVirtualMachineDeployment08.approved.json
@@ -144,7 +144,6 @@
         "storageProfile": {
           "osDisk": {
             "blobUri": "",
-            "osState": "",
             "osType": "Linux"
           }
         }
diff --git a/builder/azure/common/template/template.go b/builder/azure/common/template/template.go
index c7892e2cf..60e51b35d 100644
--- a/builder/azure/common/template/template.go
+++ b/builder/azure/common/template/template.go
@@ -35,6 +35,24 @@ type Resource struct {
 	Resources  *[]Resource         `json:"resources,omitempty"`
 }
 
+type OSDiskUnion struct {
+	OsType       compute.OperatingSystemTypes      `json:"osType,omitempty"`
+	OsState      compute.OperatingSystemStateTypes `json:"osState,omitempty"`
+	BlobURI      *string                           `json:"blobUri,omitempty"`
+	Name         *string                           `json:"name,omitempty"`
+	Vhd          *compute.VirtualHardDisk          `json:"vhd,omitempty"`
+	Image        *compute.VirtualHardDisk          `json:"image,omitempty"`
+	Caching      compute.CachingTypes              `json:"caching,omitempty"`
+	CreateOption compute.DiskCreateOptionTypes     `json:"createOption,omitempty"`
+	DiskSizeGB   *int32                            `json:"diskSizeGB,omitempty"`
+}
+
+// Union of the StorageProfile and ImageStorageProfile types.
+type StorageProfileUnion struct {
+	ImageReference *compute.ImageReference `json:"imageReference,omitempty"`
+	OsDisk         *OSDiskUnion            `json:"osDisk,omitempty"`
+}
+
 /////////////////////////////////////////////////
 // Template > Resource > Properties
 type Properties struct {
@@ -50,7 +68,8 @@ type Properties struct {
 	OsProfile                    *compute.OSProfile                  `json:"osProfile,omitempty"`
 	PublicIPAllocatedMethod      *network.IPAllocationMethod         `json:"publicIPAllocationMethod,omitempty"`
 	Sku                          *Sku                                `json:"sku,omitempty"`
-	StorageProfile               *compute.StorageProfile             `json:"storageProfile,omitempty"`
+	//StorageProfile3              *compute.StorageProfile             `json:"storageProfile,omitempty"`
+	StorageProfile               *StorageProfileUnion                `json:"storageProfile,omitempty"`
 	Subnets                      *[]network.Subnet                   `json:"subnets,omitempty"`
 	TenantId                     *string                             `json:"tenantId,omitempty"`
 	Value                        *string                             `json:"value,omitempty"`
diff --git a/builder/azure/common/template/template_builder.go b/builder/azure/common/template/template_builder.go
index c0a531863..d93094ff4 100644
--- a/builder/azure/common/template/template_builder.go
+++ b/builder/azure/common/template/template_builder.go
@@ -122,11 +122,11 @@ func (s *TemplateBuilder) SetManagedDiskUrl(managedDiskImageName, location, blob
 		Name:       to.StringPtr(managedDiskImageName),
 		Location:   to.StringPtr(location),
 		Properties: &Properties{
-			StorageProfile: &compute.StorageProfile{
-				OsDisk: &compute.OSDisk{
+			StorageProfile: &StorageProfileUnion{
+				OsDisk: &OSDiskUnion{
 					OsType:  s.osType,
-					OsState: to.StringPtr(fmt.Sprintf("%s", osState)),
-					BlobUri: to.StringPtr(blobUri),
+					OsState: osState,
+					BlobURI: to.StringPtr(blobUri),
 				},
 			},
 		},