model_3 / train.py
Junnyfirst34's picture
Create train.py
31899f4 verified
import torch
import torch.nn as nn
class SensorEncoder(nn.Module):
def __init__(self):
super().__init__()
self.fc = nn.Linear(10, 32)
def forward(self, x):
return self.fc(x)
def main():
model = SensorEncoder()
dummy_input = torch.randn(1, 10)
output = model(dummy_input)
print(output)
if __name__ == "__main__":
main()