main.py

import torch

from TorchNet_impl import TensorNet
from train_eval import run_experiment, prepare_data_qm9, extract_max_z


loaders = prepare_data_qm9(reduced_set=True)

mean = torch.tensor([data.y for data in loaders[0].dataset]).mean()
std = torch.tensor([data.y for data in loaders[0].dataset]).std()

model = TensorNet(
    hidden_dims=[64] * 4,
    max_z=extract_max_z(),
    rbf_dim=64,
    equivariance_class='SO(3)',
    predict_type='scalar',
    mean=mean,
    std=std
)


run_experiment(
    model=model,
    model_name='TensorNet',
    train_loader=loaders[0],
    val_loader=loaders[1],
    test_loader=loaders[2],
    n_epochs=100,
    patience=15
)