Logic
18 lines
550 B
Python
18 lines
550 B
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
|
|
class LinearConditionProjector(nn.Module):
|
|
"""Projects per-step visual+state conditioning to the head conditioning width."""
|
|
|
|
def __init__(self, input_dim: int, output_dim: int, bias: bool = True):
|
|
super().__init__()
|
|
self.input_dim = int(input_dim)
|
|
self.output_dim = int(output_dim)
|
|
self.linear = nn.Linear(self.input_dim, self.output_dim, bias=bias)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.linear(x)
|