fix imports

CHANGELOG.md

@@ -0,0 +1,13 @@
+# Changelog
+
+All notable changes to this project will be documented in this file.
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+### Added
+- Added multitask decoder taxonomy. ([#36](https://github.com/neuro-galaxy/torch_brain/pull/36))
+
+### Changed
+- Update workflow to use ubuntu-latest instances from github actions. ([#36](httpps://github.com/neuro-galaxy/torch_brain/pull/36))
+
+### Fixed

setup.py

@@ -11,8 +11,8 @@
     packages=find_packages() + find_namespace_packages(include=["hydra_plugins.*"]),
     include_package_data=True,
     install_requires=[
-        "temporaldata @ git+https://github.com/neuro-galaxy/temporaldata@main#egg=temporaldata",
-        "brainsets @ git+https://github.com/neuro-galaxy/brainsets@main#egg=brainsets",
+        "temporaldata @ git+https://github.com/neuro-galaxy/temporaldata@main#egg=temporaldata-0.1.1",
+        "brainsets @ git+https://github.com/neuro-galaxy/brainsets@main#egg=brainsets-0.1.0",
         "torch==2.2.0",
         "einops~=0.6.0",
         # "setuptools~=60.2.0",

torch_brain/nn/__init__.py

@@ -8,8 +8,9 @@
 from .perceiver_rotary import PerceiverRotary
 
 # readout layers
-from .loss import compute_loss_or_metric
+from .loss import compute_loss_or_metric, OutputType
 from .multitask_readout import (
     MultitaskReadout,
     prepare_for_multitask_readout,
+    Decoder,
 )

torch_brain/nn/loss.py

@@ -3,7 +3,14 @@
 
 from torchmetrics import R2Score
 
-from brainsets.taxonomy import OutputType
+from brainsets.taxonomy import StringIntEnum
+
+
+class OutputType(StringIntEnum):
+    CONTINUOUS = 0
+    BINARY = 1
+    MULTILABEL = 2
+    MULTINOMIAL = 3
 
 
 def compute_loss_or_metric(

torch_brain/nn/multitask_readout.py

@@ -7,21 +7,13 @@
 
 from brainsets.taxonomy import StringIntEnum, Task
 from torch_brain.data.collate import collate, chain, track_batch
-from torch_brain.nn import compute_loss_or_metric
-
+from torch_brain.nn import compute_loss_or_metric, OutputType
 
 from typing import Dict, List, Tuple, Optional, Union, Any
 
 from pydantic.dataclasses import dataclass
 
 
-class OutputType(StringIntEnum):
-    CONTINUOUS = 0
-    BINARY = 1
-    MULTILABEL = 2
-    MULTINOMIAL = 3
-
-
 class Decoder(StringIntEnum):
     NA = 0
     # Classic BCI outputs.