SplatAnything / src /models /layers /layer_scale.py
ritianyu's picture
Publish preview-first SplatAnything demo for ZeroGPU
dc09dcc verified
Raw History Blame Contribute Delete
511 Bytes
# Modified for SplatAnything: inference-only release.
from typing import Union
import torch
from torch import Tensor
from torch import nn
class LayerScale(nn.Module):
def __init__(self, dim: int, init_values: Union[float, Tensor]=1e-05, inplace: bool=False) -> None:
super().__init__()
self.inplace = inplace
self.gamma = nn.Parameter(init_values * torch.ones(dim))
def forward(self, x: Tensor) -> Tensor:
return x.mul_(self.gamma) if self.inplace else x * self.gamma