def get_weight(
        self,
        location_type: WeightLocationType,
        layer: LayerIndex = None,
        normalize_dim: Dimension | None = None,
        device: torch.device | None = None,
    ) -> torch.Tensor:
        """Returns the specified weights, with shape checking and optional normalization.

        Tensors returned by this method are not cloned, so please be sure not to perform in-place
        edits on them!
        """
        assert (
            location_type in weight_shape_by_location_type
        ), f"location_type_str {location_type} not found"
        weight = self._get_weight_helper(
            location_type=location_type, layer=layer, device=device or self.device
        )

        if normalize_dim is not None:
            weight = nn.functional.normalize(
                weight,
                dim=get_dimension_index_of_weight_location_type(location_type, normalize_dim),
            )
        weight_shape_spec = weight_shape_by_location_type[location_type]
        expected_shape = self.get_shape_from_spec(weight_shape_spec)
        assert (
            weight.shape == expected_shape
        ), f"Expected shape {expected_shape} but got {weight.shape}"
        # We don't want to return tensors that have gradients enabled, so we detach. Ideally we'd
        # also clone since we don't want callers to inadvertently edit the weights, but doing so
        # uses a lot of memory, so instead we just ask politely in the docstring.
        return weight.detach()
Returns the specified weights, with shape checking and optional normalization.

        Tensors returned by this method are not cloned, so please be sure not to perform in-place
        edits on them!
        
neuron_explainer/models/model_context.py