diff --git a/README.md b/README.md
index 1b7a78d..d267633 100644
--- a/README.md
+++ b/README.md
@@ -199,6 +199,27 @@ print(output)
 
 ### Firefly
 
+Exploration into the Firefly algorithm (a generalized version of particle swarm optimization) in Pytorch. In particular interested in hybrid <a href="https://academic.oup.com/jcde/article/9/2/706/6566441">firefly + genetic algorithms</a>, or ones that are <a href="https://www.sciencedirect.com/science/article/abs/pii/S0957417423005298">gender-based</a>. This code was adapted from lucidrains.
+
+```python
+from swarms_torch.firefly import FireflyOptimizer
+from torch import Tensor
+
+
+def rosenbrock(x: Tensor) -> Tensor:
+    return (
+        100 * (x[..., 1:] - x[..., :-1] ** 2) ** 2 + (1 - x[..., :-1]) ** 2
+    ).sum(dim=-1)
+
+
+if __name__ == "__main__":
+    optimizer = FireflyOptimizer(cost_function=rosenbrock)
+    optimizer.optimize()
+    best_solution = optimizer.get_best_solution()
+    print(f"Best solution: {best_solution}")
+
+```
+
 
 
 
@@ -239,4 +260,31 @@ Help us accelerate our backlog by supporting us financially! Note, we're an open
 <a href="https://polar.sh/kyegomez"><img src="https://polar.sh/embed/fund-our-backlog.svg?org=kyegomez" /></a>
 
 # License
-MIT
\ No newline at end of file
+MIT
+
+
+## Citations
+
+```bibtex
+@article{Yang2018WhyTF,
+    title   = {Why the Firefly Algorithm Works?},
+    author  = {Xin-She Yang and Xingshi He},
+    journal = {ArXiv},
+    year    = {2018},
+    volume  = {abs/1806.01632},
+    url     = {https://api.semanticscholar.org/CorpusID:46940737}
+}
+```
+
+```bibtex
+@article{article,
+    author  = {El-Shorbagy, M. and Elrefaey, Adel},
+    year    = {2022},
+    month   = {04},
+    pages   = {706-730},
+    title   = {A hybrid genetic-firefly algorithm for engineering design problems},
+    volume  = {Journal of Computational Design and Engineering, Volume 9},
+    journal = {Journal of Computational Design and Engineering},
+    doi     = {10.1093/jcde/qwac013}
+}
+```
\ No newline at end of file
diff --git a/docs/swarms/firefly.md b/docs/swarms/firefly.md
new file mode 100644
index 0000000..23762b3
--- /dev/null
+++ b/docs/swarms/firefly.md
@@ -0,0 +1,157 @@
+# FireflyOptimizer
+
+```python
+class FireflyOptimizer(cost_function, steps=5000, species=4, population_size=1000, dimensions=10, lower_bound=-4.0, upper_bound=4.0, mix_species_every=25, beta0=2.0, gamma=1.0, alpha=0.1, alpha_decay=0.995, use_genetic_algorithm=False, breed_every=10, tournament_size=100, num_children=500, use_cuda=True, verbose=True)
+```
+
+The `FireflyOptimizer` class implements the Firefly Algorithm to minimize a given objective function. It simulates the flashing behavior of fireflies to explore the search space efficiently.
+
+## Parameters
+
+- **cost_function** (callable):  
+  The objective function to minimize. Should accept a `torch.Tensor` and return a `torch.Tensor` of costs.
+
+- **steps** (int, optional):  
+  Number of optimization steps. Default: `5000`.
+
+- **species** (int, optional):  
+  Number of species in the population. Default: `4`.
+
+- **population_size** (int, optional):  
+  Number of fireflies in each species. Default: `1000`.
+
+- **dimensions** (int, optional):  
+  Dimensionality of the search space. Default: `10`.
+
+- **lower_bound** (float, optional):  
+  Lower bound of the search space. Default: `-4.0`.
+
+- **upper_bound** (float, optional):  
+  Upper bound of the search space. Default: `4.0`.
+
+- **mix_species_every** (int, optional):  
+  Interval (in steps) to mix species. Default: `25`.
+
+- **beta0** (float, optional):  
+  Base attractiveness coefficient. Default: `2.0`.
+
+- **gamma** (float, optional):  
+  Light absorption coefficient controlling intensity decay. Default: `1.0`.
+
+- **alpha** (float, optional):  
+  Randomness scaling factor. Default: `0.1`.
+
+- **alpha_decay** (float, optional):  
+  Decay rate of `alpha` per step. Default: `0.995`.
+
+- **use_genetic_algorithm** (bool, optional):  
+  Whether to include genetic algorithm operations. Default: `False`.
+
+- **breed_every** (int, optional):  
+  Steps between breeding operations when using genetic algorithm. Default: `10`.
+
+- **tournament_size** (int, optional):  
+  Number of participants in each tournament selection. Default: `100`.
+
+- **num_children** (int, optional):  
+  Number of offspring produced during breeding. Default: `500`.
+
+- **use_cuda** (bool, optional):  
+  Use CUDA for computations if available. Default: `True`.
+
+- **verbose** (bool, optional):  
+  Print progress messages during optimization. Default: `True`.
+
+## Attributes
+
+| Attribute          | Type            | Description                                            |
+|--------------------|-----------------|--------------------------------------------------------|
+| `fireflies`        | `torch.Tensor`  | Positions of the fireflies in the search space.        |
+| `device`           | `torch.device`  | Device used for computations (`cpu` or `cuda`).        |
+| `current_alpha`    | `float`         | Current value of `alpha` during optimization.          |
+
+## Methods
+
+### `optimize()`
+
+Runs the optimization loop for the specified number of steps.
+
+**Example:**
+
+```python
+optimizer.optimize()
+```
+
+### `get_best_solution()`
+
+Retrieves the best solution found by the optimizer.
+
+**Returns:**
+
+- **best_firefly** (`torch.Tensor`):  
+  The best solution vector found.
+
+**Example:**
+
+```python
+best_solution = optimizer.get_best_solution()
+print(f"Best solution: {best_solution}")
+```
+
+### `generate()`
+
+Generates a new set of fireflies, reinitializing their positions.
+
+**Returns:**
+
+- **fireflies** (`torch.Tensor`):  
+  The new set of fireflies.
+
+**Example:**
+
+```python
+optimizer.generate()
+```
+
+### `reset()`
+
+Resets the optimizer to its initial state, including `alpha` and firefly positions.
+
+**Example:**
+
+```python
+optimizer.reset()
+```
+
+---
+
+**Note:** The Firefly Algorithm is inspired by the flashing behavior of fireflies and is suitable for continuous optimization problems. This implementation allows for customization and includes optional genetic algorithm operations for enhanced performance.
+
+**Example Usage:**
+
+```python
+from swarms_torch.firefly import FireflyOptimizer
+from torch import Tensor
+
+
+def rosenbrock(x: Tensor) -> Tensor:
+    return (
+        100 * (x[..., 1:] - x[..., :-1] ** 2) ** 2 + (1 - x[..., :-1]) ** 2
+    ).sum(dim=-1)
+
+
+if __name__ == "__main__":
+    optimizer = FireflyOptimizer(
+        cost_function=rosenbrock,
+        steps=100,
+        species=10,
+        population_size=100,
+        dimensions=10,
+        lower_bound=-4,
+        upper_bound=4,
+        # Many more parameters can be set, see the documentation for more details
+    )
+    optimizer.optimize()
+    best_solution = optimizer.get_best_solution()
+    print(f"Best solution: {best_solution}")
+```
\ No newline at end of file
diff --git a/examples/fire_fly_example.py b/examples/fire_fly_example.py
new file mode 100644
index 0000000..4443d7c
--- /dev/null
+++ b/examples/fire_fly_example.py
@@ -0,0 +1,24 @@
+from swarms_torch.firefly import FireflyOptimizer
+from torch import Tensor
+
+
+def rosenbrock(x: Tensor) -> Tensor:
+    return (
+        100 * (x[..., 1:] - x[..., :-1] ** 2) ** 2 + (1 - x[..., :-1]) ** 2
+    ).sum(dim=-1)
+
+
+if __name__ == "__main__":
+    optimizer = FireflyOptimizer(
+        cost_function=rosenbrock,
+        steps=100,
+        species=10,
+        population_size=100,
+        dimensions=10,
+        lower_bound=-4,
+        upper_bound=4,
+        # Many more parameters can be set, see the documentation for more details
+    )
+    optimizer.optimize()
+    best_solution = optimizer.get_best_solution()
+    print(f"Best solution: {best_solution}")
diff --git a/swarms_torch/__init__.py b/swarms_torch/__init__.py
index 7ad5ec7..c2caab9 100644
--- a/swarms_torch/__init__.py
+++ b/swarms_torch/__init__.py
@@ -12,6 +12,7 @@
     Particle,
     TransformerParticleSwarmOptimization,
 )
+from swarms_torch.firefly import FireflyOptimizer
 from swarms_torch.structs import *  # noqa
 
 __all__ = [
@@ -28,4 +29,5 @@
     "TransformerParticleSwarmOptimization",
     "HivemindSwarm",
     "MixtureOfMambas",
+    "FireflyOptimizer",
 ]
diff --git a/swarms_torch/drone_swarm.py b/swarms_torch/drone_swarm.py
index e832a79..85768d4 100644
--- a/swarms_torch/drone_swarm.py
+++ b/swarms_torch/drone_swarm.py
@@ -309,10 +309,12 @@ def forward(
         final_neighborhood_embedding = self.neighbor_mlp(obs_neighbors)
         return final_neighborhood_embedding
 
+
 @dataclass
 class SwarmMultiHeadAttentionEncoder(nn.Module):
     dim: int
 
+
 @dataclass
 class QuadSingleHeadAttentionEncoderSim2Real(SwarmMultiHeadAttentionEncoder):
     obs_space: int
diff --git a/swarms_torch/firefly.py b/swarms_torch/firefly.py
index f09dad0..b5cc063 100644
--- a/swarms_torch/firefly.py
+++ b/swarms_torch/firefly.py
@@ -90,13 +90,22 @@ def __init__(
         self.verbose = verbose
 
         # Additional initializations
-        assert self.tournament_size <= self.population_size, "Tournament size must be less than or equal to population size."
-        assert self.num_children <= self.population_size, "Number of children must be less than or equal to population size."
-        
-        self.device = torch.device('cuda' if torch.cuda.is_available() and self.use_cuda else 'cpu')
+        assert (
+            self.tournament_size <= self.population_size
+        ), "Tournament size must be less than or equal to population size."
+        assert (
+            self.num_children <= self.population_size
+        ), "Number of children must be less than or equal to population size."
+
+        self.device = torch.device(
+            "cuda" if torch.cuda.is_available() and self.use_cuda else "cpu"
+        )
 
         # Initialize fireflies
-        self.fireflies = torch.zeros((self.species, self.population_size, self.dimensions), device=self.device).uniform_(self.lower_bound, self.upper_bound)
+        self.fireflies = torch.zeros(
+            (self.species, self.population_size, self.dimensions),
+            device=self.device,
+        ).uniform_(self.lower_bound, self.upper_bound)
 
         # Initialize alpha (in case we need to reset)
         self.current_alpha = self.alpha
@@ -109,24 +118,34 @@ def optimize(self) -> None:
             costs = self.cost_function(self.fireflies)
 
             if self.verbose:
-                logger.info(f'Step {step}: Minimum cost {costs.amin():.5f}')
+                logger.info(f"Step {step}: Minimum cost {costs.amin():.5f}")
 
             # Fireflies with lower light intensity (high cost) move towards higher intensity (lower cost)
-            move_mask = einx.greater('s i, s j -> s i j', costs, costs)
+            move_mask = einx.greater("s i, s j -> s i j", costs, costs)
 
             # Get vectors of fireflies to one another
-            delta_positions = einx.subtract('s j d, s i d -> s i j d', self.fireflies, self.fireflies)
+            delta_positions = einx.subtract(
+                "s j d, s i d -> s i j d", self.fireflies, self.fireflies
+            )
 
             distance = delta_positions.norm(dim=-1)
 
-            betas = self.beta0 * torch.exp(-self.gamma * distance ** 2)
+            betas = self.beta0 * torch.exp(-self.gamma * distance**2)
 
             # Calculate movements
-            attraction = einx.multiply('s i j, s i j d -> s i j d', move_mask * betas, delta_positions)
-            random_walk = self.current_alpha * (torch.rand_like(self.fireflies) - 0.5) * (self.upper_bound - self.lower_bound)
+            attraction = einx.multiply(
+                "s i j, s i j d -> s i j d", move_mask * betas, delta_positions
+            )
+            random_walk = (
+                self.current_alpha
+                * (torch.rand_like(self.fireflies) - 0.5)
+                * (self.upper_bound - self.lower_bound)
+            )
 
             # Move the fireflies
-            self.fireflies += einx.sum('s i j d -> s i d', attraction) + random_walk
+            self.fireflies += (
+                einx.sum("s i j d -> s i d", attraction) + random_walk
+            )
 
             self.fireflies.clamp_(min=self.lower_bound, max=self.upper_bound)
 
@@ -138,8 +157,12 @@ def optimize(self) -> None:
                 midpoint = self.population_size // 2
                 fireflies_a = self.fireflies[:, :midpoint]
                 fireflies_b = self.fireflies[:, midpoint:]
-                rotated_fireflies_b = torch.roll(fireflies_b, shifts=1, dims=(0,))
-                self.fireflies = torch.cat((fireflies_a, rotated_fireflies_b), dim=1)
+                rotated_fireflies_b = torch.roll(
+                    fireflies_b, shifts=1, dims=(0,)
+                )
+                self.fireflies = torch.cat(
+                    (fireflies_a, rotated_fireflies_b), dim=1
+                )
 
             # Genetic algorithm operations
             if self.use_genetic_algorithm and (step % self.breed_every) == 0:
@@ -156,14 +179,23 @@ def _genetic_operations(self, costs: Tensor) -> None:
         """
         fitness = 1.0 / costs
 
-        batch_randperm = torch.randn((self.species, self.num_children, self.population_size), device=self.device).argsort(dim=-1)
-        tournament_indices = batch_randperm[..., :self.tournament_size]
+        batch_randperm = torch.randn(
+            (self.species, self.num_children, self.population_size),
+            device=self.device,
+        ).argsort(dim=-1)
+        tournament_indices = batch_randperm[..., : self.tournament_size]
 
-        tournament_participants = einx.get_at('s [p], s c t -> s c t', fitness, tournament_indices)
+        tournament_participants = einx.get_at(
+            "s [p], s c t -> s c t", fitness, tournament_indices
+        )
         winners_per_tournament = tournament_participants.topk(2, dim=-1).indices
 
         # Breed the top two winners of each tournament
-        parent1, parent2 = einx.get_at('s [p] d, s c parents -> parents s c d', self.fireflies, winners_per_tournament)
+        parent1, parent2 = einx.get_at(
+            "s [p] d, s c parents -> parents s c d",
+            self.fireflies,
+            winners_per_tournament,
+        )
 
         # Uniform crossover
         crossover_mask = torch.rand_like(parent1) < 0.5
@@ -171,9 +203,13 @@ def _genetic_operations(self, costs: Tensor) -> None:
 
         # Sort the fireflies by cost and replace the worst performing with the new children
         _, sorted_indices = costs.sort(dim=-1)
-        sorted_fireflies = einx.get_at('s [p] d, s sorted -> s sorted d', self.fireflies, sorted_indices)
+        sorted_fireflies = einx.get_at(
+            "s [p] d, s sorted -> s sorted d", self.fireflies, sorted_indices
+        )
 
-        self.fireflies = torch.cat((sorted_fireflies[:, :-self.num_children], children), dim=1)
+        self.fireflies = torch.cat(
+            (sorted_fireflies[:, : -self.num_children], children), dim=1
+        )
 
     def get_best_solution(self) -> Tensor:
         """
@@ -184,12 +220,12 @@ def get_best_solution(self) -> Tensor:
         Tensor
             The best solution vector.
         """
-        fireflies_flat = einx.rearrange('s p d -> (s p) d', self.fireflies)
+        fireflies_flat = einx.rearrange("s p d -> (s p) d", self.fireflies)
         costs = self.cost_function(fireflies_flat)
         sorted_costs, sorted_indices = costs.sort(dim=-1)
         best_firefly = fireflies_flat[sorted_indices[0]]
         best_cost = sorted_costs[0]
-        logger.info(f'Best solution found with cost {best_cost:.5f}')
+        logger.info(f"Best solution found with cost {best_cost:.5f}")
         return best_firefly
 
     def generate(self) -> Tensor:
@@ -201,7 +237,10 @@ def generate(self) -> Tensor:
         Tensor
             The new set of fireflies.
         """
-        self.fireflies = torch.zeros((self.species, self.population_size, self.dimensions), device=self.device).uniform_(self.lower_bound, self.upper_bound)
+        self.fireflies = torch.zeros(
+            (self.species, self.population_size, self.dimensions),
+            device=self.device,
+        ).uniform_(self.lower_bound, self.upper_bound)
         self.current_alpha = self.alpha
         return self.fireflies
 
@@ -212,6 +251,7 @@ def reset(self) -> None:
         self.generate()
         self.current_alpha = self.alpha
 
+
 # Example usage:
 
 # def rosenbrock(x: Tensor) -> Tensor:
diff --git a/swarms_torch/swarmalators/swarmalator_base.py b/swarms_torch/swarmalators/swarmalator_base.py
index ce43b21..a01a024 100644
--- a/swarms_torch/swarmalators/swarmalator_base.py
+++ b/swarms_torch/swarmalators/swarmalator_base.py
@@ -34,9 +34,7 @@ def function_for_sigma(
 
     # Define dynamics for sigma based on our assumptions
     d_sigma = (
-        gamma * interaction_sum
-        + epsilon_a * sigma_i
-        - epsilon_r * (sigma_i**3)
+        gamma * interaction_sum + epsilon_a * sigma_i - epsilon_r * (sigma_i**3)
     )
     return d_sigma
 
diff --git a/swarms_torch/swarmalators/swarmalator_transformer.py b/swarms_torch/swarmalators/swarmalator_transformer.py
index bce24c6..003f0e0 100644
--- a/swarms_torch/swarmalators/swarmalator_transformer.py
+++ b/swarms_torch/swarmalators/swarmalator_transformer.py
@@ -1,6 +1,7 @@
 """
 Swarmalators with transformer models, SUPER EXPERIMENTAL, NEEDS WORK
 """
+
 import torch
 from torch import nn