internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: microsoft/torchgeo

Issue:
Chesapeake7 and Chesapeake13 can't have the same root directory
```
ds13 = Chesapeake13(root="data/", download=True)
ds7 = Chesapeake7(root="data/", download=True)
```

will not throw any errors, but when I sample from ds7, then RasterDataset will use the `Baywide_13Class_20132014.tif` file


Discussion:
This should be really easy to fix, we just need to add something like:
```python
filename_glob = "Baywide_13Class_20132014.tif"
```
to each dataset.

Resolving patch:
diff --git a/torchgeo/datasets/chesapeake.py b/torchgeo/datasets/chesapeake.py
--- a/torchgeo/datasets/chesapeake.py
+++ b/torchgeo/datasets/chesapeake.py
@@ -44,8 +44,6 @@ class Chesapeake(RasterDataset, abc.ABC):
       <https://chesapeakeconservancy.org/wp-content/uploads/2017/01/Chesapeake_Conservancy_Accuracy_Assessment_Methodology.pdf>`_
     """
 
-    # TODO: this shouldn't be needed, but .tif.ovr file is getting picked up
-    filename_glob = "*.tif"
     is_image = False
 
     @property
@@ -164,6 +162,7 @@ class Chesapeake7(Chesapeake):
 
     base_folder = "BAYWIDE"
     filename = "Baywide_7class_20132014.tif"
+    filename_glob = filename
     zipfile = "Baywide_7Class_20132014.zip"
     md5 = "61a4e948fb2551840b6557ef195c2084"
 
@@ -210,6 +209,7 @@ class Chesapeake13(Chesapeake):
 
     base_folder = "BAYWIDE"
     filename = "Baywide_13Class_20132014.tif"
+    filename_glob = filename
     zipfile = "Baywide_13Class_20132014.zip"
     md5 = "7e51118923c91e80e6e268156d25a4b9"
 
@@ -219,6 +219,7 @@ class ChesapeakeDC(Chesapeake):
 
     base_folder = "DC"
     filename = os.path.join("DC_11001", "DC_11001.img")
+    filename_glob = filename
     zipfile = "DC_11001.zip"
     md5 = "ed06ba7570d2955e8857d7d846c53b06"
 
@@ -228,6 +229,7 @@ class ChesapeakeDE(Chesapeake):
 
     base_folder = "DE"
     filename = "DE_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_DE_STATEWIDE.zip"
     md5 = "5e12eff3b6950c01092c7e480b38e544"
 
@@ -245,6 +247,7 @@ class ChesapeakeMD(Chesapeake):
 
     base_folder = "MD"
     filename = "MD_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_MD_STATEWIDE.zip"
     md5 = "40c7cd697a887f2ffdb601b5c114e567"
 
@@ -262,6 +265,7 @@ class ChesapeakeNY(Chesapeake):
 
     base_folder = "NY"
     filename = "NY_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_NY_STATEWIDE.zip"
     md5 = "1100078c526616454ef2e508affda915"
 
@@ -271,6 +275,7 @@ class ChesapeakePA(Chesapeake):
 
     base_folder = "PA"
     filename = "PA_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_PA_STATEWIDE.zip"
     md5 = "20a2a857c527a4dbadd6beed8b47e5ab"
 
@@ -288,6 +293,7 @@ class ChesapeakeVA(Chesapeake):
 
     base_folder = "VA"
     filename = "CIC2014_VA_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_VA_STATEWIDE.zip"
     md5 = "6f2c97deaf73bb3e1ea9b21bd7a3fc8e"
 
@@ -297,6 +303,7 @@ class ChesapeakeWV(Chesapeake):
 
     base_folder = "WV"
     filename = "WV_STATEWIDE.tif"
+    filename_glob = filename
     zipfile = "_WV_STATEWIDE.zip"
     md5 = "350621ea293651fbc557a1c3e3c64cc3"
 


Test patch:
Repository: microsoft/torchgeo

Issue:
Datasets that subclass from RasterDataset will have their custom cmaps overridden
If I have:

```
CustomDataset(RasterDataset):
  
  filename_glob = "my_tiff_that_has_a_weird_embedded_colormap.tif"
  cmap = { # some sane colors
    0: (...),
    ... 
  }
```

then the constructor of RasterDataset will override the custom cmap. I think a custom cmap shouldn't be overridden.

Fixed in #417


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -348,10 +348,11 @@ def __init__(
                 try:
                     with rasterio.open(filepath) as src:
                         # See if file has a color map
-                        try:
-                            self.cmap = src.colormap(1)
-                        except ValueError:
-                            pass
+                        if len(self.cmap) == 0:
+                            try:
+                                self.cmap = src.colormap(1)
+                            except ValueError:
+                                pass
 
                         if crs is None:
                             crs = src.crs


Test patch:
Repository: microsoft/torchgeo

Issue:
CDL dataset error
When calling `ds = CDL(root, download=True)`, the following error occurs: `rasterio._err.CPLE_AppDefinedError: The transformation is already "north up" or a transformation between pixel/line and georeferenced coordinates cannot be computed for ./data/cdl/2020_30m_cdls.tif.ovr. There is no affine transformation and no GCPs. Specify transformation option SRC_METHOD=NO_GEOTRANSFORM to bypass this check.` 
I think this occurs because `filename_glob = "*_30m_cdls.*"` also picks up the other files such as `_30m_cdls.tif.ovr` which are part of the dataset.


Discussion:
Does changing `filename_glob = "*_30m_cdls.tif"` fix it?
Yes, but the documentation of `RasterDataset`states that file extension should not be included.
We should probably include it in this case, we had to do the same thing for CDL unfortunately. See #139 for a more general discussion on whether or not it's possible to avoid this hack.

Resolving patch:
diff --git a/torchgeo/datasets/cdl.py b/torchgeo/datasets/cdl.py
--- a/torchgeo/datasets/cdl.py
+++ b/torchgeo/datasets/cdl.py
@@ -31,7 +31,7 @@ class CDL(RasterDataset):
     * https://www.nass.usda.gov/Research_and_Science/Cropland/sarsfaqs2.php#Section1_14.0
     """  # noqa: E501
 
-    filename_glob = "*_30m_cdls.*"
+    filename_glob = "*_30m_cdls.tif"
     filename_regex = r"""
         ^(?P<date>\d+)
         _30m_cdls\..*$


Test patch:
Repository: microsoft/torchgeo

Issue:
datasets.SEN12MS - dB values are cast to int32
When loading samples from the SEN12MS dataset, the Sentinel-1 dB values (floats ranging from around -30 to 0) are cast to int32, discarding a lot of important information. 

[This line](https://github.com/microsoft/torchgeo/blob/main/torchgeo/datasets/sen12ms.py#L269) seems to be the culprit.

Fixing this would require either breaking the current behaviour where S1 and S2 imagery are stacked into a single tensor, or casting everything to float32 (not sure if this is okay for S2 data)


Discussion:
I vote for casting everything to float32. I think PyTorch will automatically do this for us, so all you have to do is remove the cast to int32. Want to open a PR?

Resolving patch:
diff --git a/torchgeo/datasets/sen12ms.py b/torchgeo/datasets/sen12ms.py
--- a/torchgeo/datasets/sen12ms.py
+++ b/torchgeo/datasets/sen12ms.py
@@ -266,7 +266,9 @@ def _load_raster(self, filename: str, source: str) -> Tensor:
                 "{}_{}_{}_{}_{}".format(*parts),
             )
         ) as f:
-            array = f.read().astype(np.int32)
+            array = f.read()
+            if array.dtype == np.uint16:
+                array = array.astype(np.int32)
             tensor = torch.from_numpy(array)
             return tensor
 


Test patch:
Repository: microsoft/torchgeo

Issue:
Docs failing due to pytorch-lightning deprecation
Docs are failing due to the recent deprecation of `core/lightning.py` in favor of `core/module.py` in pytorch-lightning. Refer PyTorchLightning/pytorch-lightning#12740. 





Discussion:

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -107,10 +107,10 @@
     "matplotlib": ("https://matplotlib.org/stable/", None),
     "numpy": ("https://numpy.org/doc/stable/", None),
     "python": ("https://docs.python.org/3", None),
-    "pytorch-lightning": ("https://pytorch-lightning.readthedocs.io/en/latest/", None),
-    "rasterio": ("https://rasterio.readthedocs.io/en/latest/", None),
-    "rtree": ("https://rtree.readthedocs.io/en/latest/", None),
-    "segmentation_models_pytorch": ("https://smp.readthedocs.io/en/latest/", None),
+    "pytorch-lightning": ("https://pytorch-lightning.readthedocs.io/en/stable/", None),
+    "rasterio": ("https://rasterio.readthedocs.io/en/stable/", None),
+    "rtree": ("https://rtree.readthedocs.io/en/stable/", None),
+    "segmentation_models_pytorch": ("https://smp.readthedocs.io/en/stable/", None),
     "torch": ("https://pytorch.org/docs/stable", None),
     "torchvision": ("https://pytorch.org/vision/stable", None),
 }


Test patch:
Repository: microsoft/torchgeo

Issue:
Landcover.ai download link
It looks like the landcover.ai domain name expired (breaking our download). The dataset is now hosted at https://landcover.ai.linuxpolska.com/.


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/landcoverai.py b/torchgeo/datasets/landcoverai.py
--- a/torchgeo/datasets/landcoverai.py
+++ b/torchgeo/datasets/landcoverai.py
@@ -60,7 +60,7 @@ class LandCoverAI(VisionDataset):
          the train/val/test split
     """
 
-    url = "https://landcover.ai/download/landcover.ai.v1.zip"
+    url = "https://landcover.ai.linuxpolska.com/download/landcover.ai.v1.zip"
     filename = "landcover.ai.v1.zip"
     md5 = "3268c89070e8734b4e91d531c0617e03"
     sha256 = "15ee4ca9e3fd187957addfa8f0d74ac31bc928a966f76926e11b3c33ea76daa1"


Test patch:
Repository: microsoft/torchgeo

Issue:
Update README
Before blog: update verbiage and add pics from blog/paper. 

After blog: add link to blog.


Discussion:
Note: for any images we add, we'll need to make sure that they get added to the sdist so that they display properly on PyPI.

Resolving patch:
diff --git a/torchgeo/__init__.py b/torchgeo/__init__.py
--- a/torchgeo/__init__.py
+++ b/torchgeo/__init__.py
@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
-"""TorchGeo: datasets, transforms, and models for geospatial data.
+"""TorchGeo: datasets, samplers, transforms, and pre-trained models for geospatial data.
 
 This library is part of the `PyTorch <http://pytorch.org/>`_ project. PyTorch is an open
 source machine learning framework.


Test patch:
Repository: microsoft/torchgeo

Issue:
VectorDataset incorrect sample shape 
### Description

Querying a sample from a VectorDataset object returns (sometimes?) a mask with incorrect shape.

This is caused by not rounding the calculated width and height pixel values to the nearest integer, and just taking the `int(width)` and `int(height)` values instead for `masks = rasterio.features.rasterize` in [torchgeo.datasets.geo.py:614](https://github.com/microsoft/torchgeo/blob/b41b6d68f99edfbe4d82b31f6f9b16cbd33f485e/torchgeo/datasets/geo.py#L614).

### Steps to reproduce

Run the following code to reproduce the error:
```
from torchgeo.datasets import CanadianBuildingFootprints
from torchgeo.samplers import RandomGeoSampler

cbf_dataset = CanadianBuildingFootprints(download=True)  # or any other VectorDataset-based dataset

size = 512

sampler = RandomGeoSampler(cbf_dataset, size=size, length=100)  # or any other sampler
query = next(iter(sampler))

sample = cbf_dataset[query]

print(sample["mask"].shape)
assert sample["mask"].shape[0] == size, sample["mask"].shape
assert sample["mask"].shape[1] == size, sample["mask"].shape
```

### Version

0.3.0


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -611,12 +611,12 @@ def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
         )
         if shapes:
             masks = rasterio.features.rasterize(
-                shapes, out_shape=(int(height), int(width)), transform=transform
+                shapes, out_shape=(round(height), round(width)), transform=transform
             )
         else:
             # If no features are found in this query, return an empty mask
             # with the default fill value and dtype used by rasterize
-            masks = np.zeros((int(height), int(width)), dtype=np.uint8)
+            masks = np.zeros((round(height), round(width)), dtype=np.uint8)
 
         sample = {"mask": torch.tensor(masks), "crs": self.crs, "bbox": query}
 


Test patch:
diff --git a/tests/datasets/test_geo.py b/tests/datasets/test_geo.py
--- a/tests/datasets/test_geo.py
+++ b/tests/datasets/test_geo.py
@@ -223,7 +223,7 @@ def test_getitem(self, dataset: CustomVectorDataset) -> None:
     def test_empty_shapes(self, dataset: CustomVectorDataset) -> None:
         query = BoundingBox(1.1, 1.9, 1.1, 1.9, 0, 0)
         x = dataset[query]
-        assert torch.equal(x["mask"], torch.zeros(7, 7, dtype=torch.uint8))
+        assert torch.equal(x["mask"], torch.zeros(8, 8, dtype=torch.uint8))
 
     def test_invalid_query(self, dataset: CustomVectorDataset) -> None:
         query = BoundingBox(3, 3, 3, 3, 0, 0)
Repository: microsoft/torchgeo

Issue:
Add CSVLogger to train.py
### Summary

Our current `train.py` script only logs to TensorBoard. We should add a CSV logger.

### Rationale

In order to determine the optimal hyperparams in a hyperparam sweep, we need an easy way of finding the final val accuracy of each run. Currently, the only way to get this is to convert the TensorBoard logger output to a CSV. If we log directly to CSV this becomes much easier, especially for those who don't use TensorBoard.

### Implementation

_No response_

### Alternatives

_No response_

### Additional information

_No response_


Discussion:

Resolving patch:
diff --git a/train.py b/train.py
--- a/train.py
+++ b/train.py
@@ -170,6 +170,7 @@ def main(conf: DictConfig) -> None:
     # Setup trainer
     ######################################
     tb_logger = pl_loggers.TensorBoardLogger(conf.program.log_dir, name=experiment_name)
+    csv_logger = pl_loggers.CSVLogger(conf.program.log_dir, name=experiment_name)
 
     if isinstance(task, ObjectDetectionTask):
         monitor_metric = "val_map"
@@ -192,7 +193,7 @@ def main(conf: DictConfig) -> None:
     trainer_args = cast(Dict[str, Any], OmegaConf.to_object(conf.trainer))
 
     trainer_args["callbacks"] = [checkpoint_callback, early_stopping_callback]
-    trainer_args["logger"] = tb_logger
+    trainer_args["logger"] = [tb_logger, csv_logger]
     trainer_args["default_root_dir"] = experiment_dir
     trainer = pl.Trainer(**trainer_args)
 


Test patch:
Repository: microsoft/torchgeo

Issue:
EuroSAT Dataset URL not being secure anymore
### Description

As `url = "https://madm.dfki.de/files/sentinel/EuroSATallBands.zip"` is not secure anymore, downloading the EuroSat dataset is blocked by default by certify.

We might need to either change the zip source or remove the `https`

### Steps to reproduce

Load the EuroSAT dataset
```python
from torchgeo.datasets import EuroSAT
eurosat_train = EuroSAT(download=True)
```

### Version

0.4.0


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/eurosat.py b/torchgeo/datasets/eurosat.py
--- a/torchgeo/datasets/eurosat.py
+++ b/torchgeo/datasets/eurosat.py
@@ -51,7 +51,9 @@ class EuroSAT(NonGeoClassificationDataset):
     * https://ieeexplore.ieee.org/document/8519248
     """
 
-    url = "https://madm.dfki.de/files/sentinel/EuroSATallBands.zip"  # 2.0 GB download
+    # TODO: Change to https after https://github.com/phelber/EuroSAT/issues/10 is
+    # resolved
+    url = "http://madm.dfki.de/files/sentinel/EuroSATallBands.zip"  # 2.0 GB download
     filename = "EuroSATallBands.zip"
     md5 = "5ac12b3b2557aa56e1826e981e8e200e"
 
@@ -104,6 +106,8 @@ class EuroSAT(NonGeoClassificationDataset):
 
     BAND_SETS = {"all": all_band_names, "rgb": rgb_bands}
 
+    # TODO: reset checksum to False after https://github.com/phelber/EuroSAT/issues/10
+    # is resolved
     def __init__(
         self,
         root: str = "data",
@@ -111,7 +115,7 @@ def __init__(
         bands: Sequence[str] = BAND_SETS["all"],
         transforms: Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]] = None,
         download: bool = False,
-        checksum: bool = False,
+        checksum: bool = True,
     ) -> None:
         """Initialize a new EuroSAT dataset instance.
 


Test patch:
Repository: microsoft/torchgeo

Issue:
EuroSAT Dataset URL not being secure anymore
### Description

As `url = "https://madm.dfki.de/files/sentinel/EuroSATallBands.zip"` is not secure anymore, downloading the EuroSat dataset is blocked by default by certify.

We might need to either change the zip source or remove the `https`

### Steps to reproduce

Load the EuroSAT dataset
```python
from torchgeo.datasets import EuroSAT
eurosat_train = EuroSAT(download=True)
```

### Version

0.4.0


Discussion:
https://github.com/microsoft/torchgeo/pull/1065
Thanks @julien-blanchon -- we opened an issue in the EuroSAT repo last week (https://github.com/phelber/EuroSAT/issues/10) and were hoping that they fix the expired certificate ASAP. It looks like we will need to switch to http though.

Resolving patch:
diff --git a/torchgeo/datasets/eurosat.py b/torchgeo/datasets/eurosat.py
--- a/torchgeo/datasets/eurosat.py
+++ b/torchgeo/datasets/eurosat.py
@@ -51,9 +51,7 @@ class EuroSAT(NonGeoClassificationDataset):
     * https://ieeexplore.ieee.org/document/8519248
     """
 
-    # TODO: Change to https after https://github.com/phelber/EuroSAT/issues/10 is
-    # resolved
-    url = "http://madm.dfki.de/files/sentinel/EuroSATallBands.zip"  # 2.0 GB download
+    url = "https://huggingface.co/datasets/torchgeo/eurosat/resolve/main/EuroSATallBands.zip"  # noqa: E501
     filename = "EuroSATallBands.zip"
     md5 = "5ac12b3b2557aa56e1826e981e8e200e"
 
@@ -106,8 +104,6 @@ class EuroSAT(NonGeoClassificationDataset):
 
     BAND_SETS = {"all": all_band_names, "rgb": rgb_bands}
 
-    # TODO: reset checksum to False after https://github.com/phelber/EuroSAT/issues/10
-    # is resolved
     def __init__(
         self,
         root: str = "data",
@@ -115,7 +111,7 @@ def __init__(
         bands: Sequence[str] = BAND_SETS["all"],
         transforms: Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]] = None,
         download: bool = False,
-        checksum: bool = True,
+        checksum: bool = False,
     ) -> None:
         """Initialize a new EuroSAT dataset instance.
 


Test patch:
Repository: microsoft/torchgeo

Issue:
Remove date from NLCD glob/regex
### Description

We should remove the date from the NLCD glob/regex. NLCD products are regenerated every time a new release comes out, meaning the name (and unfortunately the hash) will change.

### Steps to reproduce

This came to me in a dream. 

### Version

0.5.0dev


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/nlcd.py b/torchgeo/datasets/nlcd.py
--- a/torchgeo/datasets/nlcd.py
+++ b/torchgeo/datasets/nlcd.py
@@ -65,11 +65,11 @@ class NLCD(RasterDataset):
     .. versionadded:: 0.5
     """  # noqa: E501
 
-    filename_glob = "nlcd_*_land_cover_l48_20210604.img"
+    filename_glob = "nlcd_*_land_cover_l48_*.img"
     filename_regex = (
         r"nlcd_(?P<date>\d{4})_land_cover_l48_(?P<publication_date>\d{8})\.img"
     )
-    zipfile_glob = "nlcd_*_land_cover_l48_20210604.zip"
+    zipfile_glob = "nlcd_*_land_cover_l48_*.zip"
     date_format = "%Y"
     is_image = False
 
@@ -188,7 +188,7 @@ def _verify(self) -> None:
         # Check if the extracted files already exist
         exists = []
         for year in self.years:
-            filename_year = self.filename_glob.replace("*", str(year))
+            filename_year = self.filename_glob.replace("*", str(year), 1)
             pathname = os.path.join(self.root, "**", filename_year)
             if glob.glob(pathname, recursive=True):
                 exists.append(True)
@@ -201,7 +201,7 @@ def _verify(self) -> None:
         # Check if the zip files have already been downloaded
         exists = []
         for year in self.years:
-            zipfile_year = self.zipfile_glob.replace("*", str(year))
+            zipfile_year = self.zipfile_glob.replace("*", str(year), 1)
             pathname = os.path.join(self.root, "**", zipfile_year)
             if glob.glob(pathname, recursive=True):
                 exists.append(True)
@@ -236,7 +236,7 @@ def _download(self) -> None:
     def _extract(self) -> None:
         """Extract the dataset."""
         for year in self.years:
-            zipfile_name = self.zipfile_glob.replace("*", str(year))
+            zipfile_name = self.zipfile_glob.replace("*", str(year), 1)
             pathname = os.path.join(self.root, "**", zipfile_name)
             extract_archive(glob.glob(pathname, recursive=True)[0], self.root)
 


Test patch:
Repository: microsoft/torchgeo

Issue:
USAVars Augmentation maps to 0
### Description

In the USAVars Datamodule, the default augmentation from NonGeoDatamodule is used. However, the dataset returns uint8 data, and it comes out of the augmentation still as uint8. This means you get an error when trying to train but also that your input images are just all zeros.

### Steps to reproduce
```
dm = USAVarsDataModule(root="path/to/usa_vars", batch_size=16)
dm.setup("fit")
dl = dm.train_dataloader()
batch = next(iter(dl))
aug_batch = dm.aug(batch)
print(aug_batch["image"].max())
```

### Version

'0.5.0.dev0'


Discussion:

Resolving patch:
diff --git a/torchgeo/datasets/usavars.py b/torchgeo/datasets/usavars.py
--- a/torchgeo/datasets/usavars.py
+++ b/torchgeo/datasets/usavars.py
@@ -200,7 +200,7 @@ def _load_image(self, path: str) -> Tensor:
         """
         with rasterio.open(path) as f:
             array: "np.typing.NDArray[np.int_]" = f.read()
-            tensor = torch.from_numpy(array)
+            tensor = torch.from_numpy(array).float()
             return tensor
 
     def _verify(self) -> None:


Test patch:
Repository: microsoft/torchgeo

Issue:
Add class weights parameter to classification/detection trainers
### Summary

#1221 adds a class_weights parameter to the SemanticSegmentation trainer. We should do the same for other class-based trainers.

### Rationale

Class imbalance is common in geospatial datasets.

### Implementation

Copy implementation from #1221

### Alternatives

_No response_

### Additional information

_No response_


Discussion:

Resolving patch:
diff --git a/torchgeo/trainers/classification.py b/torchgeo/trainers/classification.py
--- a/torchgeo/trainers/classification.py
+++ b/torchgeo/trainers/classification.py
@@ -38,6 +38,7 @@ def __init__(
         in_channels: int = 3,
         num_classes: int = 1000,
         loss: str = "ce",
+        class_weights: Optional[Tensor] = None,
         lr: float = 1e-3,
         patience: int = 10,
         freeze_backbone: bool = False,
@@ -53,6 +54,8 @@ def __init__(
             in_channels: Number of input channels to model.
             num_classes: Number of prediction classes.
             loss: One of 'ce', 'bce', 'jaccard', or 'focal'.
+            class_weights: Optional rescaling weight given to each
+                class and used with 'ce' loss.
             lr: Learning rate for optimizer.
             patience: Patience for learning rate scheduler.
             freeze_backbone: Freeze the backbone network to linear probe
@@ -62,7 +65,7 @@ def __init__(
            *classification_model* was renamed to *model*.
 
         .. versionadded:: 0.5
-           The *freeze_backbone* parameter.
+           The *class_weights* and *freeze_backbone* parameters.
 
         .. versionchanged:: 0.5
            *learning_rate* and *learning_rate_schedule_patience* were renamed to
@@ -78,7 +81,9 @@ def configure_losses(self) -> None:
         """
         loss: str = self.hparams["loss"]
         if loss == "ce":
-            self.criterion: nn.Module = nn.CrossEntropyLoss()
+            self.criterion: nn.Module = nn.CrossEntropyLoss(
+                weight=self.hparams["class_weights"]
+            )
         elif loss == "bce":
             self.criterion = nn.BCEWithLogitsLoss()
         elif loss == "jaccard":


Test patch: