Gram Matrix for Style

Implement a function to compute the Gram matrix for neural style transfer, which is a crucial component in capturing the style of an image. This task involves understanding how to represent the style of an image using feature maps extracted from a Convolutional Neural Network (CNN).

The Gram matrix is used to capture the style by computing correlations between feature maps, where the feature map $F$ is a 3D tensor with shape $(C, H \times W)$, representing the output of a CNN layer with $C$ channels, height $H$, and width $W$. The Gram matrix $G$ is a 2D tensor that represents the correlations between these feature maps.

Here are the steps to compute the Gram matrix:

1. Extract the feature map $F$ from a CNN layer.
2. Compute the correlations between the feature maps. The Gram matrix $G$ can be computed using the formula $G_{ij} = \sum_k F_{ik} F_{jk}$, which represents the dot product of the $i^{th}$ and $j^{th}$ feature maps.

The style loss $L_{style}$ is then computed as the squared difference between the Gram matrix of the content image and the Gram matrix of the style image: $L_{style} = \|G_{content} - G_{style}\|^2$.

This technique is widely used in image generation and editing tasks.