diff --git a/Nbs/01_layers.ipynb b/Nbs/01_layers.ipynb
index 3b95728..d52275f 100644
--- a/Nbs/01_layers.ipynb
+++ b/Nbs/01_layers.ipynb
@@ -2,261 +2,261 @@
  "cells": [
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "# Layers\n",
     "\n",
     "> Basic layers for constructor."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 1,
-   "source": [
-    "#hide\n",
-    "from nbdev.showdoc import show_doc\n",
-    "from IPython.display import Markdown, display"
-   ],
-   "outputs": [],
    "metadata": {
     "ExecuteTime": {
      "end_time": "2021-08-11T16:34:43.007417Z",
      "start_time": "2021-08-11T16:34:42.829843Z"
     }
-   }
+   },
+   "outputs": [],
+   "source": [
+    "#hide\n",
+    "from nbdev.showdoc import show_doc\n",
+    "from IPython.display import Markdown, display"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
-   "source": [
-    "#hide\n",
-    "import torch.nn as nn\n",
-    "import torch\n",
-    "from collections import OrderedDict"
-   ],
-   "outputs": [],
    "metadata": {
     "ExecuteTime": {
      "end_time": "2021-08-11T16:34:45.654722Z",
      "start_time": "2021-08-11T16:34:45.085125Z"
     }
-   }
+   },
+   "outputs": [],
+   "source": [
+    "#hide\n",
+    "import torch.nn as nn\n",
+    "import torch\n",
+    "from collections import OrderedDict"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "# hide\n",
     "def print_doc(func_name):\n",
     "    doc = show_doc(func_name, title_level=4, disp=False)\n",
     "    display(Markdown(doc))"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Flatten layer"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
-   "source": [
-    "#hide\n",
-    "from model_constructor.layers import Flatten, noop, Noop"
-   ],
-   "outputs": [],
    "metadata": {
     "ExecuteTime": {
      "end_time": "2021-08-11T16:34:45.996818Z",
      "start_time": "2021-08-11T16:34:45.982654Z"
     }
-   }
+   },
+   "outputs": [],
+   "source": [
+    "#hide\n",
+    "from model_constructor.layers import Flatten, noop, Noop"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 5,
-   "source": [
-    "#hide_input\n",
-    "# show_doc(Flatten, title_level=4)\n",
-    "# flatten_doc = show_doc(Flatten, title_level=4, disp=False)\n",
-    "# Markdown(flatten_doc)\n",
-    "\n",
-    "print_doc(Flatten)"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T16:34:54.998866Z",
+     "start_time": "2021-08-11T16:34:54.984142Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
        "<h4 id=\"Flatten\" class=\"doc_header\"><code>class</code> <code>Flatten</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L11\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>Flatten</code>() :: `Module`\n",
        "\n",
        "flat x to vector"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T16:34:54.998866Z",
-     "start_time": "2021-08-11T16:34:54.984142Z"
-    }
-   }
+   "source": [
+    "#hide_input\n",
+    "# show_doc(Flatten, title_level=4)\n",
+    "# flatten_doc = show_doc(Flatten, title_level=4, disp=False)\n",
+    "# Markdown(flatten_doc)\n",
+    "\n",
+    "print_doc(Flatten)"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## Noop - dummy func and module."
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 6,
-   "source": [
-    "#hide_input\n",
-    "# show_doc(noop)\n",
-    "# doc = show_doc(noop, disp=False)\n",
-    "# display(Markdown(doc))\n",
-    "print_doc(noop)"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:55.119562Z",
+     "start_time": "2021-08-11T15:44:55.103484Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
        "<h4 id=\"noop\" class=\"doc_header\"><code>noop</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L20\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>noop</code>(**`x`**)\n",
        "\n",
        "Dummy func. Return input"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:55.119562Z",
-     "start_time": "2021-08-11T15:44:55.103484Z"
-    }
-   }
+   "source": [
+    "#hide_input\n",
+    "# show_doc(noop)\n",
+    "# doc = show_doc(noop, disp=False)\n",
+    "# display(Markdown(doc))\n",
+    "print_doc(noop)"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 7,
-   "source": [
-    "#hide_input\n",
-    "# show_doc(Noop, title_level=4)\n",
-    "# doc = show_doc(Noop, title_level=4, disp=False)\n",
-    "# Markdown(doc)\n",
-    "print_doc(Noop)"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T16:35:14.225068Z",
+     "start_time": "2021-08-11T16:35:14.208899Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
        "<h4 id=\"Noop\" class=\"doc_header\"><code>class</code> <code>Noop</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L25\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>Noop</code>() :: `Module`\n",
        "\n",
        "Dummy module"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T16:35:14.225068Z",
-     "start_time": "2021-08-11T16:35:14.208899Z"
-    }
-   }
+   "source": [
+    "#hide_input\n",
+    "# show_doc(Noop, title_level=4)\n",
+    "# doc = show_doc(Noop, title_level=4, disp=False)\n",
+    "# Markdown(doc)\n",
+    "print_doc(Noop)"
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## ConvLayer - nn.module"
-   ],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 8,
-   "source": [
-    "#hide\n",
-    "from model_constructor.layers import ConvLayer"
-   ],
-   "outputs": [],
    "metadata": {
     "ExecuteTime": {
      "end_time": "2021-08-11T16:35:21.435560Z",
      "start_time": "2021-08-11T16:35:21.430089Z"
     }
-   }
+   },
+   "outputs": [],
+   "source": [
+    "#hide\n",
+    "from model_constructor.layers import ConvLayer"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 9,
-   "source": [
-    "# hide_input\n",
-    "# show_doc(ConvLayer, title_level=4)\n",
-    "# doc = show_doc(ConvLayer, title_level=4, disp=False)\n",
-    "# Markdown(doc)\n",
-    "print_doc(ConvLayer)"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T16:35:29.238261Z",
+     "start_time": "2021-08-11T16:35:29.229345Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
        "<h4 id=\"ConvLayer\" class=\"doc_header\"><code>class</code> <code>ConvLayer</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L37\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>ConvLayer</code>(**`ni`**, **`nf`**, **`ks`**=*`3`*, **`stride`**=*`1`*, **`act`**=*`True`*, **`act_fn`**=*`ReLU(inplace=True)`*, **`bn_layer`**=*`True`*, **`bn_1st`**=*`True`*, **`zero_bn`**=*`False`*, **`padding`**=*`None`*, **`bias`**=*`False`*, **`groups`**=*`1`*, **\\*\\*`kwargs`**) :: `Sequential`\n",
        "\n",
        "Basic conv layers block"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T16:35:29.238261Z",
-     "start_time": "2021-08-11T16:35:29.229345Z"
-    }
-   }
+   "source": [
+    "# hide_input\n",
+    "# show_doc(ConvLayer, title_level=4)\n",
+    "# doc = show_doc(ConvLayer, title_level=4, disp=False)\n",
+    "# Markdown(doc)\n",
+    "print_doc(ConvLayer)"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 10,
-   "source": [
-    "#collapse_output\n",
-    "conv_layer = ConvLayer(32, 64, act=False)\n",
-    "conv_layer"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:56.242450Z",
+     "start_time": "2021-08-11T15:44:56.231392Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "ConvLayer(\n",
@@ -265,28 +265,28 @@
        ")"
       ]
      },
+     "execution_count": 10,
      "metadata": {},
-     "execution_count": 10
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:56.242450Z",
-     "start_time": "2021-08-11T15:44:56.231392Z"
-    }
-   }
+   "source": [
+    "#collapse_output\n",
+    "conv_layer = ConvLayer(32, 64, act=False)\n",
+    "conv_layer"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 11,
-   "source": [
-    "#collapse_output\n",
-    "conv_layer = ConvLayer(32, 64, bn_layer=False)\n",
-    "conv_layer"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:56.462424Z",
+     "start_time": "2021-08-11T15:44:56.450423Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "ConvLayer(\n",
@@ -295,28 +295,28 @@
        ")"
       ]
      },
+     "execution_count": 11,
      "metadata": {},
-     "execution_count": 11
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:56.462424Z",
-     "start_time": "2021-08-11T15:44:56.450423Z"
-    }
-   }
+   "source": [
+    "#collapse_output\n",
+    "conv_layer = ConvLayer(32, 64, bn_layer=False)\n",
+    "conv_layer"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 12,
-   "source": [
-    "#collapse_output\n",
-    "conv_layer = ConvLayer(32, 64, bn_1st=True)\n",
-    "conv_layer"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:56.667542Z",
+     "start_time": "2021-08-11T15:44:56.660962Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "ConvLayer(\n",
@@ -326,28 +326,28 @@
        ")"
       ]
      },
+     "execution_count": 12,
      "metadata": {},
-     "execution_count": 12
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:56.667542Z",
-     "start_time": "2021-08-11T15:44:56.660962Z"
-    }
-   }
+   "source": [
+    "#collapse_output\n",
+    "conv_layer = ConvLayer(32, 64, bn_1st=True)\n",
+    "conv_layer"
+   ]
   },
   {
    "cell_type": "code",
    "execution_count": 13,
-   "source": [
-    "#collapse_output\n",
-    "conv_layer = ConvLayer(32, 64, bn_1st=True, act_fn=nn.LeakyReLU())\n",
-    "conv_layer"
-   ],
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:56.905543Z",
+     "start_time": "2021-08-11T15:44:56.900817Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "ConvLayer(\n",
@@ -357,20 +357,27 @@
        ")"
       ]
      },
+     "execution_count": 13,
      "metadata": {},
-     "execution_count": 13
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:56.905543Z",
-     "start_time": "2021-08-11T15:44:56.900817Z"
-    }
-   }
-  },
+   "source": [
+    "#collapse_output\n",
+    "conv_layer = ConvLayer(32, 64, bn_1st=True, act_fn=nn.LeakyReLU())\n",
+    "conv_layer"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 14,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:57.154044Z",
+     "start_time": "2021-08-11T15:44:57.115266Z"
+    }
+   },
+   "outputs": [],
    "source": [
     "#hide\n",
     "bs = 8\n",
@@ -378,317 +385,903 @@
     "y = conv_layer(xb)\n",
     "y.shape\n",
     "assert y.shape == torch.Size([bs, 64, 32, 32])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## SimpleSelfAttention"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "SA module from mxresnet at fastai."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# hide\n",
+    "from model_constructor.layers import conv1d, SimpleSelfAttention"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"conv1d\" class=\"doc_header\"><code>conv1d</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L65\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>conv1d</code>(**`ni`**:`int`, **`no`**:`int`, **`ks`**:`int`=*`1`*, **`stride`**:`int`=*`1`*, **`padding`**:`int`=*`0`*, **`bias`**:`bool`=*`False`*)\n",
+       "\n",
+       "Create and initialize a `nn.Conv1d` layer with spectral normalization."
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
    ],
+   "source": [
+    "#hide_input\n",
+    "# doc = show_doc(conv1d, title_level=4, disp=False)\n",
+    "# Markdown(doc)\n",
+    "print_doc(conv1d)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"SimpleSelfAttention\" class=\"doc_header\"><code>class</code> <code>SimpleSelfAttention</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L74\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>SimpleSelfAttention</code>(**`n_in`**:`int`, **`ks`**=*`1`*, **`sym`**=*`False`*) :: `Module`\n",
+       "\n",
+       "SimpleSelfAttention module.  # noqa W291\n",
+       "Adapted from SelfAttention layer at  \n",
+       "https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py  \n",
+       "Inspired by https://arxiv.org/pdf/1805.08318.pdf  "
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# doc = show_doc(SimpleSelfAttention, title_level=4, disp=False)\n",
+    "# Markdown(doc)\n",
+    "print_doc(SimpleSelfAttention)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## SEModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
    "outputs": [],
+   "source": [
+    "#hide\n",
+    "from model_constructor.layers import SEModule"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"SEModule\" class=\"doc_header\"><code>class</code> <code>SEModule</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L154\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>SEModule</code>(**`channels`**, **`reduction`**=*`16`*, **`rd_channels`**=*`None`*, **`rd_max`**=*`False`*, **`se_layer`**=*`Linear`*, **`act_fn`**=*`ReLU(inplace=True)`*, **`use_bias`**=*`True`*, **`gate`**=*`Sigmoid`*) :: `Module`\n",
+       "\n",
+       "se block"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# hide_input\n",
+    "# Markdown(show_doc(SEBlock, title_level=4))\n",
+    "print_doc(SEModule)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:57.154044Z",
-     "start_time": "2021-08-11T15:44:57.115266Z"
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModule(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (fc_reduce): Linear(in_features=128, out_features=8, bias=True)\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (fc_expand): Linear(in_features=8, out_features=128, bias=True)\n",
+       "    (se_gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModule(128)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModule(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (fc_reduce): Linear(in_features=128, out_features=4, bias=True)\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (fc_expand): Linear(in_features=4, out_features=128, bias=True)\n",
+       "    (se_gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModule(128, reduction=32)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
     }
-   }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModule(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (fc_reduce): Linear(in_features=128, out_features=32, bias=True)\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (fc_expand): Linear(in_features=32, out_features=128, bias=True)\n",
+       "    (se_gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModule(128, rd_channels=32)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModule(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (fc_reduce): Linear(in_features=128, out_features=32, bias=True)\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (fc_expand): Linear(in_features=32, out_features=128, bias=True)\n",
+       "    (se_gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModule(128, reduction=4, rd_channels=16, rd_max=True)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 32,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
-    "## SimpleSelfAttention"
+    "## SEModuleConv"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/markdown": [
+       "<h4 id=\"SEModuleConv\" class=\"doc_header\"><code>class</code> <code>SEModuleConv</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L189\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "\n",
+       "> <code>SEModuleConv</code>(**`channels`**, **`reduction`**=*`16`*, **`rd_channels`**=*`None`*, **`rd_max`**=*`False`*, **`se_layer`**=*`Conv2d`*, **`act_fn`**=*`ReLU(inplace=True)`*, **`use_bias`**=*`True`*, **`gate`**=*`Sigmoid`*) :: `Module`\n",
+       "\n",
+       "se block with conv on excitation"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# hide_input\n",
+    "from model_constructor.layers import SEModuleConv\n",
+    "print_doc(SEModuleConv)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModuleConv(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 35,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModuleConv(128)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModuleConv(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (conv_reduce): Conv2d(128, 4, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (conv_expand): Conv2d(4, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModuleConv(128, reduction=32)\n",
+    "se_block"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
+   ],
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "SEModuleConv(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (conv_reduce): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (conv_expand): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (gate): Sigmoid()\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 39,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
    ],
-   "metadata": {}
-  },
-  {
-   "cell_type": "markdown",
    "source": [
-    "SA module from mxresnet at fastai."
-   ],
-   "metadata": {}
+    "#collapse_output\n",
+    "se_block = SEModuleConv(128, rd_channels=32)\n",
+    "se_block"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
-   "source": [
-    "# hide\n",
-    "from model_constructor.layers import conv1d, SimpleSelfAttention"
+   "execution_count": 40,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
+    }
    ],
-   "outputs": [],
-   "metadata": {}
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
-   "source": [
-    "#hide_input\n",
-    "# doc = show_doc(conv1d, title_level=4, disp=False)\n",
-    "# Markdown(doc)\n",
-    "print_doc(conv1d)"
-   ],
+   "execution_count": 41,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
       "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
-      "text/markdown": [
-       "<h4 id=\"conv1d\" class=\"doc_header\"><code>conv1d</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L65\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
-       "\n",
-       "> <code>conv1d</code>(**`ni`**:`int`, **`no`**:`int`, **`ks`**:`int`=*`1`*, **`stride`**:`int`=*`1`*, **`padding`**:`int`=*`0`*, **`bias`**:`bool`=*`False`*)\n",
-       "\n",
-       "Create and initialize a `nn.Conv1d` layer with spectral normalization."
+       "SEModuleConv(\n",
+       "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
+       "  (excitation): Sequential(\n",
+       "    (conv_reduce): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (se_act): ReLU(inplace=True)\n",
+       "    (conv_expand): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1))\n",
+       "    (gate): Sigmoid()\n",
+       "  )\n",
+       ")"
       ]
      },
-     "metadata": {}
+     "execution_count": 41,
+     "metadata": {},
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {}
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEModuleConv(128, reduction=4, rd_channels=16, rd_max=True)\n",
+    "se_block"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
-   "source": [
-    "# doc = show_doc(SimpleSelfAttention, title_level=4, disp=False)\n",
-    "# Markdown(doc)\n",
-    "print_doc(SimpleSelfAttention)"
-   ],
+   "execution_count": 42,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "display_data",
-     "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
-      "text/markdown": [
-       "<h4 id=\"SimpleSelfAttention\" class=\"doc_header\"><code>class</code> <code>SimpleSelfAttention</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L74\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
-       "\n",
-       "> <code>SimpleSelfAttention</code>(**`n_in`**:`int`, **`ks`**=*`1`*, **`sym`**=*`False`*) :: `Module`\n",
-       "\n",
-       "SimpleSelfAttention module.  \n",
-       "Adapted from SelfAttention layer at  \n",
-       "https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py  \n",
-       "Inspired by https://arxiv.org/pdf/1805.08318.pdf  "
-      ]
-     },
-     "metadata": {}
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "torch.Size([8, 128, 32, 32])\n"
+     ]
     }
    ],
-   "metadata": {}
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
-    "## SE Block"
-   ],
-   "metadata": {}
+    "## SEBlock"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "First version of SE block, leaved for compatibility."
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 43,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "#hide\n",
     "from model_constructor.layers import SEBlock"
-   ],
-   "outputs": [],
-   "metadata": {}
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
-   "source": [
-    "# hide_input\n",
-    "# Markdown(show_doc(SEBlock, title_level=4))\n",
-    "print_doc(SEBlock)"
-   ],
+   "execution_count": 44,
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
        "<h4 id=\"SEBlock\" class=\"doc_header\"><code>class</code> <code>SEBlock</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L105\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>SEBlock</code>(**`c`**, **`r`**=*`16`*) :: `Module`\n",
        "\n",
        "se block"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {}
+   "source": [
+    "# hide_input\n",
+    "# Markdown(show_doc(SEBlock, title_level=4))\n",
+    "print_doc(SEBlock)"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
-   "source": [
-    "#collapse_output\n",
-    "se_block = SEBlock(128)\n",
-    "se_block"
-   ],
+   "execution_count": 45,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.343217Z",
+     "start_time": "2021-08-11T15:44:58.332669Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "SEBlock(\n",
        "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
        "  (excitation): Sequential(\n",
-       "    (fc_reduce): Linear(in_features=128, out_features=8, bias=False)\n",
+       "    (fc_reduce): Linear(in_features=128, out_features=8, bias=True)\n",
        "    (se_act): ReLU(inplace=True)\n",
-       "    (fc_expand): Linear(in_features=8, out_features=128, bias=False)\n",
+       "    (fc_expand): Linear(in_features=8, out_features=128, bias=True)\n",
        "    (sigmoid): Sigmoid()\n",
        "  )\n",
        ")"
       ]
      },
+     "execution_count": 45,
      "metadata": {},
-     "execution_count": 21
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:58.343217Z",
-     "start_time": "2021-08-11T15:44:58.332669Z"
-    }
-   }
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEBlock(128)\n",
+    "se_block"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
-   "source": [
-    "#hide\n",
-    "bs = 8\n",
-    "xb = torch.randn(bs, 128, 32, 32)\n",
-    "y = se_block(xb)\n",
-    "print(y.shape)\n",
-    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
-   ],
+   "execution_count": 46,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:58.570883Z",
+     "start_time": "2021-08-11T15:44:58.530689Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "torch.Size([8, 128, 32, 32])\n"
      ]
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:58.570883Z",
-     "start_time": "2021-08-11T15:44:58.530689Z"
-    }
-   }
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## SEBlockConv"
-   ],
-   "metadata": {}
+   ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 23,
+   "cell_type": "markdown",
+   "metadata": {},
    "source": [
-    "# hide_input\n",
-    "from model_constructor.layers import SEBlockConv\n",
-    "print_doc(SEBlockConv)"
-   ],
+    "First version of SEBlockConv, leaved for compatibility."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "metadata": {},
    "outputs": [
     {
-     "output_type": "display_data",
      "data": {
-      "text/plain": [
-       "<IPython.core.display.Markdown object>"
-      ],
       "text/markdown": [
-       "<h4 id=\"SEBlockConv\" class=\"doc_header\"><code>class</code> <code>SEBlockConv</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L128\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
+       "<h4 id=\"SEBlockConv\" class=\"doc_header\"><code>class</code> <code>SEBlockConv</code><a href=\"https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L129\" class=\"source_link\" style=\"float:right\">[source]</a></h4>\n",
        "\n",
        "> <code>SEBlockConv</code>(**`c`**, **`r`**=*`16`*) :: `Module`\n",
        "\n",
        "se block with conv on excitation"
+      ],
+      "text/plain": [
+       "<IPython.core.display.Markdown object>"
       ]
      },
-     "metadata": {}
+     "metadata": {},
+     "output_type": "display_data"
     }
    ],
-   "metadata": {}
+   "source": [
+    "# hide_input\n",
+    "from model_constructor.layers import SEBlockConv\n",
+    "print_doc(SEBlockConv)"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
-   "source": [
-    "#collapse_output\n",
-    "se_block = SEBlockConv(128)\n",
-    "se_block"
-   ],
+   "execution_count": 48,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:59.189301Z",
+     "start_time": "2021-08-11T15:44:59.177022Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "execute_result",
      "data": {
       "text/plain": [
        "SEBlockConv(\n",
        "  (squeeze): AdaptiveAvgPool2d(output_size=1)\n",
        "  (excitation): Sequential(\n",
-       "    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))\n",
        "    (se_act): ReLU(inplace=True)\n",
-       "    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
+       "    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))\n",
        "    (sigmoid): Sigmoid()\n",
        "  )\n",
        ")"
       ]
      },
+     "execution_count": 48,
      "metadata": {},
-     "execution_count": 24
+     "output_type": "execute_result"
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:59.189301Z",
-     "start_time": "2021-08-11T15:44:59.177022Z"
-    }
-   }
+   "source": [
+    "#collapse_output\n",
+    "se_block = SEBlockConv(128)\n",
+    "se_block"
+   ]
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
-   "source": [
-    "#hide\n",
-    "bs = 8\n",
-    "xb = torch.randn(bs, 128, 32, 32)\n",
-    "y = se_block(xb)\n",
-    "print(y.shape)\n",
-    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
-   ],
+   "execution_count": 49,
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2021-08-11T15:44:59.406923Z",
+     "start_time": "2021-08-11T15:44:59.362600Z"
+    }
+   },
    "outputs": [
     {
-     "output_type": "stream",
      "name": "stdout",
+     "output_type": "stream",
      "text": [
       "torch.Size([8, 128, 32, 32])\n"
      ]
     }
    ],
-   "metadata": {
-    "ExecuteTime": {
-     "end_time": "2021-08-11T15:44:59.406923Z",
-     "start_time": "2021-08-11T15:44:59.362600Z"
-    }
-   }
+   "source": [
+    "#hide\n",
+    "bs = 8\n",
+    "xb = torch.randn(bs, 128, 32, 32)\n",
+    "y = se_block(xb)\n",
+    "print(y.shape)\n",
+    "assert y.shape == torch.Size([bs, 128, 32, 32]), f\"size\""
+   ]
   },
   {
    "cell_type": "markdown",
+   "metadata": {},
    "source": [
     "## end\n",
     "model_constructor\n",
     "by ayasyrev"
-   ],
-   "metadata": {}
+   ]
   }
  ],
  "metadata": {
+  "interpreter": {
+   "hash": "460c8d17e5de1304fcc9388854d8b1e7fdd10d3c58b2d7b68fabbdff2124405d"
+  },
   "kernelspec": {
-   "name": "python3",
-   "display_name": "Python 3.9.6 64-bit ('mc_dev': conda)"
+   "display_name": "Python 3.9.6 64-bit ('mc_dev': conda)",
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -700,7 +1293,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.6"
+   "version": "3.8.10"
   },
   "toc": {
    "base_numbering": 1,
@@ -719,11 +1312,8 @@
    },
    "toc_section_display": true,
    "toc_window_display": true
-  },
-  "interpreter": {
-   "hash": "460c8d17e5de1304fcc9388854d8b1e7fdd10d3c58b2d7b68fabbdff2124405d"
   }
  },
  "nbformat": 4,
  "nbformat_minor": 2
-}
\ No newline at end of file
+}
diff --git a/docs/layers.html b/docs/layers.html
index 2e183a7..8062319 100644
--- a/docs/layers.html
+++ b/docs/layers.html
@@ -375,7 +375,7 @@ <h4 id="conv1d" class="doc_header"><code>conv1d</code><a href="https://github.co
 <div class="output_markdown rendered_html output_subarea ">
 <h4 id="SimpleSelfAttention" class="doc_header"><code>class</code> <code>SimpleSelfAttention</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L74" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>SimpleSelfAttention</code>(<strong><code>n_in</code></strong>:<code>int</code>, <strong><code>ks</code></strong>=<em><code>1</code></em>, <strong><code>sym</code></strong>=<em><code>False</code></em>) :: <code>Module</code></p>
 </blockquote>
-<p>SimpleSelfAttention module.<br>
+<p>SimpleSelfAttention module.  # noqa W291
 Adapted from SelfAttention layer at<br>
 <a href="https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py">https://github.com/fastai/fastai/blob/5c51f9eabf76853a89a9bc5741804d2ed4407e49/fastai/layers.py</a><br>
 Inspired by <a href="https://arxiv.org/pdf/1805.08318.pdf">https://arxiv.org/pdf/1805.08318.pdf</a></p>
@@ -392,7 +392,444 @@ <h4 id="SimpleSelfAttention" class="doc_header"><code>class</code> <code>SimpleS
 
 <div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
 <div class="text_cell_render border-box-sizing rendered_html">
-<h2 id="SE-Block">SE Block<a class="anchor-link" href="#SE-Block"> </a></h2>
+<h2 id="SEModule">SEModule<a class="anchor-link" href="#SEModule"> </a></h2>
+</div>
+</div>
+</div>
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+<div class="output_markdown rendered_html output_subarea ">
+<h4 id="SEModule" class="doc_header"><code>class</code> <code>SEModule</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L154" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>SEModule</code>(<strong><code>channels</code></strong>, <strong><code>reduction</code></strong>=<em><code>16</code></em>, <strong><code>rd_channels</code></strong>=<em><code>None</code></em>, <strong><code>rd_max</code></strong>=<em><code>False</code></em>, <strong><code>se_layer</code></strong>=<em><code>Linear</code></em>, <strong><code>act_fn</code></strong>=<em><code>ReLU(inplace=True)</code></em>, <strong><code>use_bias</code></strong>=<em><code>True</code></em>, <strong><code>gate</code></strong>=<em><code>Sigmoid</code></em>) :: <code>Module</code></p>
+</blockquote>
+<p>se block</p>
+
+</div>
+
+</div>
+
+</div>
+</div>
+
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModule</span><span class="p">(</span><span class="mi">128</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModule(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (fc_reduce): Linear(in_features=128, out_features=8, bias=True)
+    (se_act): ReLU(inplace=True)
+    (fc_expand): Linear(in_features=8, out_features=128, bias=True)
+    (se_gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModule</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="mi">32</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModule(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (fc_reduce): Linear(in_features=128, out_features=4, bias=True)
+    (se_act): ReLU(inplace=True)
+    (fc_expand): Linear(in_features=4, out_features=128, bias=True)
+    (se_gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModule</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">rd_channels</span><span class="o">=</span><span class="mi">32</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModule(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (fc_reduce): Linear(in_features=128, out_features=32, bias=True)
+    (se_act): ReLU(inplace=True)
+    (fc_expand): Linear(in_features=32, out_features=128, bias=True)
+    (se_gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModule</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span> <span class="n">rd_channels</span><span class="o">=</span><span class="mi">16</span><span class="p">,</span> <span class="n">rd_max</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModule(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (fc_reduce): Linear(in_features=128, out_features=32, bias=True)
+    (se_act): ReLU(inplace=True)
+    (fc_expand): Linear(in_features=32, out_features=128, bias=True)
+    (se_gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
+<div class="text_cell_render border-box-sizing rendered_html">
+<h2 id="SEModuleConv">SEModuleConv<a class="anchor-link" href="#SEModuleConv"> </a></h2>
+</div>
+</div>
+</div>
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+<div class="output_markdown rendered_html output_subarea ">
+<h4 id="SEModuleConv" class="doc_header"><code>class</code> <code>SEModuleConv</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L189" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>SEModuleConv</code>(<strong><code>channels</code></strong>, <strong><code>reduction</code></strong>=<em><code>16</code></em>, <strong><code>rd_channels</code></strong>=<em><code>None</code></em>, <strong><code>rd_max</code></strong>=<em><code>False</code></em>, <strong><code>se_layer</code></strong>=<em><code>Conv2d</code></em>, <strong><code>act_fn</code></strong>=<em><code>ReLU(inplace=True)</code></em>, <strong><code>use_bias</code></strong>=<em><code>True</code></em>, <strong><code>gate</code></strong>=<em><code>Sigmoid</code></em>) :: <code>Module</code></p>
+</blockquote>
+<p>se block with conv on excitation</p>
+
+</div>
+
+</div>
+
+</div>
+</div>
+
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModuleConv</span><span class="p">(</span><span class="mi">128</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModuleConv(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))
+    (se_act): ReLU(inplace=True)
+    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))
+    (gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModuleConv</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="mi">32</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModuleConv(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (conv_reduce): Conv2d(128, 4, kernel_size=(1, 1), stride=(1, 1))
+    (se_act): ReLU(inplace=True)
+    (conv_expand): Conv2d(4, 128, kernel_size=(1, 1), stride=(1, 1))
+    (gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModuleConv</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">rd_channels</span><span class="o">=</span><span class="mi">32</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModuleConv(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (conv_reduce): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1))
+    (se_act): ReLU(inplace=True)
+    (conv_expand): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1))
+    (gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+    {% raw %}
+    
+<div class="cell border-box-sizing code_cell rendered">
+<div class="input">
+
+<div class="inner_cell">
+    <div class="input_area">
+<div class=" highlight hl-ipython3"><pre><span></span><span class="n">se_block</span> <span class="o">=</span> <span class="n">SEModuleConv</span><span class="p">(</span><span class="mi">128</span><span class="p">,</span> <span class="n">reduction</span><span class="o">=</span><span class="mi">4</span><span class="p">,</span> <span class="n">rd_channels</span><span class="o">=</span><span class="mi">16</span><span class="p">,</span> <span class="n">rd_max</span><span class="o">=</span><span class="kc">True</span><span class="p">)</span>
+<span class="n">se_block</span>
+</pre></div>
+
+    </div>
+</div>
+</div>
+<details class="description">
+      <summary data-open="Hide Output" data-close="Show Output"></summary>
+        <summary></summary>
+        
+<div class="output_wrapper">
+<div class="output">
+
+<div class="output_area">
+
+
+
+<div class="output_text output_subarea output_execute_result">
+<pre>SEModuleConv(
+  (squeeze): AdaptiveAvgPool2d(output_size=1)
+  (excitation): Sequential(
+    (conv_reduce): Conv2d(128, 32, kernel_size=(1, 1), stride=(1, 1))
+    (se_act): ReLU(inplace=True)
+    (conv_expand): Conv2d(32, 128, kernel_size=(1, 1), stride=(1, 1))
+    (gate): Sigmoid()
+  )
+)</pre>
+</div>
+
+</div>
+
+</div>
+</div>
+
+    </details>
+</div>
+    {% endraw %}
+
+<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
+<div class="text_cell_render border-box-sizing rendered_html">
+<h2 id="SEBlock">SEBlock<a class="anchor-link" href="#SEBlock"> </a></h2>
+</div>
+</div>
+</div>
+<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
+<div class="text_cell_render border-box-sizing rendered_html">
+<p>First version of SE block, leaved for compatibility.</p>
+
 </div>
 </div>
 </div>
@@ -450,9 +887,9 @@ <h4 id="SEBlock" class="doc_header"><code>class</code> <code>SEBlock</code><a hr
 <pre>SEBlock(
   (squeeze): AdaptiveAvgPool2d(output_size=1)
   (excitation): Sequential(
-    (fc_reduce): Linear(in_features=128, out_features=8, bias=False)
+    (fc_reduce): Linear(in_features=128, out_features=8, bias=True)
     (se_act): ReLU(inplace=True)
-    (fc_expand): Linear(in_features=8, out_features=128, bias=False)
+    (fc_expand): Linear(in_features=8, out_features=128, bias=True)
     (sigmoid): Sigmoid()
   )
 )</pre>
@@ -470,6 +907,13 @@ <h4 id="SEBlock" class="doc_header"><code>class</code> <code>SEBlock</code><a hr
 <div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
 <div class="text_cell_render border-box-sizing rendered_html">
 <h2 id="SEBlockConv">SEBlockConv<a class="anchor-link" href="#SEBlockConv"> </a></h2>
+</div>
+</div>
+</div>
+<div class="cell border-box-sizing text_cell rendered"><div class="inner_cell">
+<div class="text_cell_render border-box-sizing rendered_html">
+<p>First version of SEBlockConv, leaved for compatibility.</p>
+
 </div>
 </div>
 </div>
@@ -484,7 +928,7 @@ <h2 id="SEBlockConv">SEBlockConv<a class="anchor-link" href="#SEBlockConv"> </a>
 
 
 <div class="output_markdown rendered_html output_subarea ">
-<h4 id="SEBlockConv" class="doc_header"><code>class</code> <code>SEBlockConv</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L128" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>SEBlockConv</code>(<strong><code>c</code></strong>, <strong><code>r</code></strong>=<em><code>16</code></em>) :: <code>Module</code></p>
+<h4 id="SEBlockConv" class="doc_header"><code>class</code> <code>SEBlockConv</code><a href="https://github.com/ayasyrev/model_constructor/tree/master/model_constructor/layers.py#L129" class="source_link" style="float:right">[source]</a></h4><blockquote><p><code>SEBlockConv</code>(<strong><code>c</code></strong>, <strong><code>r</code></strong>=<em><code>16</code></em>) :: <code>Module</code></p>
 </blockquote>
 <p>se block with conv on excitation</p>
 
@@ -527,9 +971,9 @@ <h4 id="SEBlockConv" class="doc_header"><code>class</code> <code>SEBlockConv</co
 <pre>SEBlockConv(
   (squeeze): AdaptiveAvgPool2d(output_size=1)
   (excitation): Sequential(
-    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
+    (conv_reduce): Conv2d(128, 8, kernel_size=(1, 1), stride=(1, 1))
     (se_act): ReLU(inplace=True)
-    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
+    (conv_expand): Conv2d(8, 128, kernel_size=(1, 1), stride=(1, 1))
     (sigmoid): Sigmoid()
   )
 )</pre>
diff --git a/model_constructor/layers.py b/model_constructor/layers.py
index 1ef3497..f9799dc 100644
--- a/model_constructor/layers.py
+++ b/model_constructor/layers.py
@@ -157,13 +157,20 @@ class SEModule(nn.Module):
     def __init__(self,
                  channels,
                  reduction=16,
+                 rd_channels=None,
+                 rd_max=False,
                  se_layer=nn.Linear,
                  act_fn=nn.ReLU(inplace=True),  # ? obj or class?
                  use_bias=True,
                  gate=nn.Sigmoid
                  ):
         super().__init__()
-        rd_channels = channels // reduction
+        reducted = channels // reduction
+        if rd_channels is None:
+            rd_channels = reducted
+        else:
+            if rd_max:
+                rd_channels = max(rd_channels, reducted)
         self.squeeze = nn.AdaptiveAvgPool2d(1)
         self.excitation = nn.Sequential(
             OrderedDict([('fc_reduce', se_layer(channels, rd_channels, bias=use_bias)),
@@ -185,6 +192,8 @@ class SEModuleConv(nn.Module):
     def __init__(self,
                  channels,
                  reduction=16,
+                 rd_channels=None,
+                 rd_max=False,
                  se_layer=nn.Conv2d,
                  act_fn=nn.ReLU(inplace=True),
                  use_bias=True,
@@ -192,7 +201,12 @@ def __init__(self,
                  ):
         super().__init__()
 #       rd_channels = math.ceil(channels//reduction/8)*8
-        rd_channels = channels // reduction
+        reducted = channels // reduction
+        if rd_channels is None:
+            rd_channels = reducted
+        else:
+            if rd_max:
+                rd_channels = max(rd_channels, reducted)
         self.squeeze = nn.AdaptiveAvgPool2d(1)
         self.excitation = nn.Sequential(
             OrderedDict([